Update 16_tensorboard.py

model was on graphic card and example_data was on cpu so was throwing error
lcong · Jan 18, 2022 · b5843ae · b5843ae
1 parent e6dd99c
commit b5843ae
Showing 1 changed file with 2 additions and 2 deletions.
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()
-    ###################################################
+    ###################################################