Merge pull request patrickloeber#13 from preyasgarg/patch-2

patrickloeber · web-flow · commit 69637721ed9f · 2023-01-05T16:22:03.000+01:00
Update 16_tensorboard.py
diff --git a/16_tensorboard.py b/16_tensorboard.py
@@ -80,7 +80,7 @@ def forward(self, x):
 optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  
 
 ############## TENSORBOARD ########################
-writer.add_graph(model, example_data.reshape(-1, 28*28))
+writer.add_graph(model, example_data.reshape(-1, 28*28).to(device))
 #writer.close()
 #sys.exit()
 ###################################################
@@ -156,4 +156,4 @@ def forward(self, x):
         preds_i = class_preds[:, i]
         writer.add_pr_curve(str(i), labels_i, preds_i, global_step=0)
         writer.close()
-    ###################################################
+    ###################################################
