add warning for not staged training model

pytorch · Apr 16, 2024 · 92071f1 · 92071f1
1 parent 0315d0b
commit 92071f1
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/torchbenchmark/util/extra_args.py b/torchbenchmark/util/extra_args.py
@@ -188,9 +188,10 @@ def apply_decoration_args(
             model.add_context(lambda: torch.cuda.amp.autocast(dtype=torch.float16))
         elif model.test == "train":
             # the model must implement staged train test
-            assert is_staged_train_test(
-                model
-            ), f"Expected model implements staged train test (forward, backward, optimizer)."
+            warnings.warn(
+                    "Usually models do not want to enable amp only in forward path, so expected "
+                    "model to have staged train support."
+                    )
             import torch
 
             model.add_context(

diff --git a/torchbenchmark/util/model.py b/torchbenchmark/util/model.py
@@ -483,8 +483,18 @@ 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 self.test == "train" and is_staged_train_test(self):
-            self.forward_contexts.append(self.amp_context)
+        if self.test == "eval":
+            self.add_context(self.amp_context)
+        elif self.test == "train":
+            if is_staged_train_test(self):
+                self.add_context(self.amp_context, TEST_STAGE.FORWARD)
+            else:
+                warnings.warn(
+                        "Usually models do not want to enable amp only in forward path, so expected "
+                        "model to have staged train support."
+                        )
+
+
 
     @property
     def pt2_compilation_time(self) -> Optional[float]: