writer.add_graph fix bug

dingguanglei · Nov 15, 2018 · d58c0c5 · d58c0c5
1 parent e4f1888
commit d58c0c5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jdit/trainer/super.py b/jdit/trainer/super.py
@@ -360,7 +360,7 @@ def graph(self, model: Union[torch.nn.Module, torch.nn.DataParallel, Model],
             res = proto_model(*input)
             self.scalars({'ParamsNum': num_params}, 1, tag="ParamsNum")
             del res
-            self.writer.add_graph(proto_model, *input)
+            self.writer.add_graph(proto_model, input)
 
     def _count_params(self, proto_model: torch.nn.Module):
         """count the total parameters of model.