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

Update 16_tensorboard.py
lcong · Jan 5, 2023 · 6963772 · 6963772
2 parents fbb6f59 + b5843ae
commit 6963772
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()
-    ###################################################
+    ###################################################