fix state dict device

Author: vmoens

test/test_collector.py

@@ -1515,7 +1515,7 @@ def create_env():
                 ].keys()
                 for k in state_dict[f"worker{worker}"]["policy_state_dict"]:
                     torch.testing.assert_close(
-                        state_dict[f"worker{worker}"]["policy_state_dict"][k],
+                        state_dict[f"worker{worker}"]["policy_state_dict"][k].cpu(),
                         policy_state_dict[k].cpu(),
                     )
 
@@ -1533,7 +1533,7 @@ def create_env():
                         AssertionError
                     ) if torch.cuda.is_available() else nullcontext():
                         torch.testing.assert_close(
-                            state_dict[f"worker{worker}"]["policy_state_dict"][k],
+                            state_dict[f"worker{worker}"]["policy_state_dict"][k].cpu(),
                             policy_state_dict[k].cpu(),
                         )
 
@@ -1546,7 +1546,7 @@ def create_env():
             for worker in range(3):
                 for k in state_dict[f"worker{worker}"]["policy_state_dict"]:
                     torch.testing.assert_close(
-                        state_dict[f"worker{worker}"]["policy_state_dict"][k],
+                        state_dict[f"worker{worker}"]["policy_state_dict"][k].cpu(),
                         policy_state_dict[k].cpu(),
                     )
         finally: