Update

[ghstack-poisoned]
pytorch · Dec 15, 2024 · 05e7d88 · 05e7d88
2 parents f01b33f + d230184
commit 05e7d88
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sota-implementations/ppo/utils_atari.py b/sota-implementations/ppo/utils_atari.py
@@ -115,7 +115,7 @@ def make_ppo_modules_pixels(proof_environment, device):
         strides=[4, 2, 1],
         device=device,
     )
-    common_cnn_output = common_cnn(torch.ones(input_shape))
+    common_cnn_output = common_cnn(torch.ones(input_shape, device=device))
     common_mlp = MLP(
         in_features=common_cnn_output.shape[-1],
         activation_class=torch.nn.ReLU,