Add loss inputs for PT functional loss

Author: vandanavk

smdebug/pytorch/hook.py

@@ -137,12 +137,15 @@ def forward_pre_hook(self, module, inputs):
             self.export_collections()
             self.exported_collections = True
 
-    def record_tensor_value(self, tensor_name: str, tensor_value: torch.Tensor) -> None:
+    def record_tensor_value(self, tensor_name: str, tensor_value: torch.Tensor, inputs=None) -> None:
         """Used for registering functional directly, such as F.mse_loss()."""
         assert isinstance(
             tensor_value, torch.Tensor
         ), f"tensor_value={tensor_value} must be torch.Tensor"
 
+        if inputs:
+            self._write_inputs(tensor_name, inputs)
+
         self._write_outputs(tensor_name, tensor_value)
 
     # This hook is invoked by trainer after running the forward pass.

tests/pytorch/test_loss.py

@@ -58,7 +58,7 @@ def create_net_and_train(out_dir, n_steps, use_loss_module=False, use_loss_funct
             loss = criterion(outputs, labels)
         if use_loss_functional:
             loss = F.cross_entropy(outputs, labels)
-            hook.record_tensor_value("nll_loss", tensor_value=loss)
+            hook.record_tensor_value("nll_loss", tensor_value=loss, inputs=[outputs, labels])
         loss.backward()
         optimizer.step()
 
@@ -78,8 +78,8 @@ def test_register_loss_functional(out_dir):
     loss_tensor = trial.tensor("nll_loss_output_0")
 
     # Capture ['nll_loss_output_0']
-    assert len(trial.tensor_names()) == 1
-    assert len(loss_coll.tensor_names) == 1
+    assert len(trial.tensor_names()) == 3
+    assert len(loss_coll.tensor_names) == 3
 
     # Loss should be logged for all the steps since passed `available_steps = range(n_steps)`
     assert len(trial.steps()) == n_steps