fix(wandb_torch): extract graph.criterion for RNN models
@@ -255,6 +255,8 @@ class TorchGraph(wandb.data_types.Graph):
if not graph.criterion_passed:
if hasattr(output[0], 'grad_fn'):
graph.criterion = output[0].grad_fn
+ elif isinstance(output[0], list) and hasattr(output[0][0], 'grad_fn'):
+ graph.criterion = output[0][0].grad_fn
return after_forward_hook
def hook_torch_modules(self, module, criterion=None, prefix=None, graph_idx=0):
@@ -449,3 +451,4 @@ class TorchGraph(wandb.data_types.Graph):
node.class_name = type(module).__name__
return node
+
