.

FilipinoGambino · Feb 8, 2024 · e79d26f · e79d26f
1 parent 6a973e3
commit e79d26f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/connectx/nns/models.py b/connectx/nns/models.py
@@ -192,7 +192,8 @@ def forward(
     ) -> Dict[str, Any]:
         logging.info(f"Beginning forward pass")
         x, available_actions_mask, subtask_embeddings = self.dict_input_layer(x)
-        logging.info(f"Getting base_model outputs {x.shape}")
+        for key,val in x.items():
+            logging.info(f"Getting base_model outputs {key}:{val.shape}")
         base_out = self.base_model(x)
         logging.info(f"Ignoring subtasks")
         if subtask_embeddings is not None: