amend

facebookresearch · Apr 5, 2024 · 627c8ff · 627c8ff
1 parent 2c48aba
commit 627c8ff
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/benchmarl/models/cnn.py b/benchmarl/models/cnn.py
@@ -130,7 +130,10 @@ def __init__(
                 device=self.device,
                 **cnn_net_kwargs,
             )
-            example_net = self.cnn._empty_net
+            if hasattr(self.cnn, "_empy_net"):
+                example_net = self.cnn._empty_net
+            else:
+                example_net = self.cnn.agent_networks[0]
         else:
             self.cnn = nn.ModuleList(
                 [