[fix] Use torch.inference_mode inplace of torch.no_grad (#3188)

aimhubio · Jul 14, 2024 · 756f41a · 756f41a
1 parent dd246b3
commit 756f41a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/examples/pytorch_track.py b/examples/pytorch_track.py
@@ -12,8 +12,8 @@
 # Initialize a new Run
 aim_run = Run()
 
-# Device configuration
-device = torch.device('cpu')
+# moving model to gpu if available
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
 # Hyper parameters
 num_epochs = 5
@@ -118,8 +118,7 @@ def forward(self, x):
 
 
 # Test the model
-model.eval()
-with torch.no_grad():
+with torch.inference_mode():
     correct = 0
     total = 0
     for images, labels in test_loader:

diff --git a/examples/pytorch_track_images.py b/examples/pytorch_track_images.py
@@ -13,8 +13,8 @@
 # Initialize a new Run
 aim_run = Run()
 
-# Device configuration
-device = torch.device('cpu')
+# moving model to gpu if available
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
 # Hyper parameters
 num_epochs = 5
@@ -122,8 +122,7 @@ def forward(self, x):
 
 
 # Test the model
-model.eval()
-with torch.no_grad():
+with torch.inference_mode():
     correct = 0
     total = 0
     for images, labels in tqdm(test_loader, total=len(test_loader)):