add condition in enable_amp to fix eval not work under amp

pytorch · Apr 11, 2024 · 0315d0b · 0315d0b
1 parent 6caee05
commit 0315d0b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchbenchmark/util/model.py b/torchbenchmark/util/model.py
@@ -483,7 +483,7 @@ def enable_amp(self):
             self.amp_context = lambda: torch.cpu.amp.autocast()
         elif self.device == "cuda":
             self.amp_context = lambda: torch.cuda.amp.autocast()
-        if is_staged_train_test(self):
+        if self.test == "train" and is_staged_train_test(self):
             self.forward_contexts.append(self.amp_context)
 
     @property