Minor issue when training on grid env

[Jama Hussein Mohamud]

• 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