[BugFix] device in road_traffic scenario (#157)

Author: Schopenhauer-loves-Hegel

vmas/scenarios/road_traffic.py

@@ -745,7 +745,11 @@ def init_params(self, batch_dim, device, **kwargs):
 
         self.distances = Distances(
             agents=torch.zeros(
-                batch_dim, self.n_agents, self.n_agents, dtype=torch.float32
+                batch_dim,
+                self.n_agents,
+                self.n_agents,
+                device=device,
+                dtype=torch.float32,
             ),
             left_boundaries=torch.zeros(
                 (batch_dim, self.n_agents, 1 + 4), device=device, dtype=torch.float32