down into cpp

pytorch · Apr 24, 2024 · f244aab · f244aab
1 parent 5734811
commit f244aab
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/test_test_mnist.py b/test/test_test_mnist.py
@@ -89,8 +89,8 @@ def forward(self, x):
     return F.log_softmax(x, dim=1)
 
 device = xm.xla_device()
-# model = MNIST().to(device)
-network = Net().to(device)
+network = MNIST().to(device)
+# network = Net().to(device)
 optimizer = optim.SGD(network.parameters(), lr=learning_rate,
                       momentum=momentum)
 # loss_fn = nn.NLLLoss()