From e79d26fcad0439b63837bd494ef8b136845a0722 Mon Sep 17 00:00:00 2001 From: FilipinoGambino Date: Wed, 7 Feb 2024 20:58:35 -0600 Subject: [PATCH] . --- connectx/nns/models.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/connectx/nns/models.py b/connectx/nns/models.py index 48028f8..7a3b1cf 100644 --- 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: