Only set graph  criterion if we know what we are doing.
@@ -253,6 +253,7 @@ class TorchGraph(wandb.data_types.Graph):
graph.nodes_by_id[id(param)] = node
graph.add_node(node)
if not graph.criterion_passed:
+ if hasattr(output[0], 'grad_fn'):
graph.criterion = output[0].grad_fn
return after_forward_hook
