[BugFix] Add missing break (#2322)

pytorch · Jul 25, 2024 · 94abb50 · 94abb50
1 parent dd70b78
commit 94abb50
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/test/test_modules.py b/test/test_modules.py
@@ -863,6 +863,13 @@ def test_multiagent_mlp_init(
             share_params=share_params,
             depth=2,
         )
+        for m in mlp.modules():
+            if isinstance(m, nn.Linear):
+                assert not isinstance(m.weight, nn.Parameter)
+                assert m.weight.device == torch.device("meta")
+                break
+        else:
+            raise RuntimeError("could not find a Linear module")
         if n_agent_inputs is None:
             n_agent_inputs = 6
         td = self._get_mock_input_td(n_agents, n_agent_inputs, batch=batch)