Minor issue when training on grid env

[engmubarak48]

• Puts all tensors into the device
  change:
  

  gflownet/gflownet/envs/grid.py

  Line 181 in bcb1443

  cols = states + torch.arange(self.n_dim) * self.length

  
  To:

            cols = states + torch.arange(self.n_dim, device=self.device) * self.length