torchtnt/tests/framework (#952)

Summary: Pull Request resolved: #952 Reviewed By: galrotem, avikchaudhuri Differential Revision: D67388469 fbshipit-source-id: 6b0fce0e95c1561407c9e118b2b45d47299c8b56
pytorch · Dec 19, 2024 · 2e762a1 · 2e762a1
1 parent 2ce298f
commit 2e762a1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/framework/test_loop_utils.py b/tests/framework/test_loop_utils.py
@@ -105,7 +105,7 @@ def forward(self, x):
                 return x
 
         loss_fn = nn.CrossEntropyLoss()
-        module = torch.export.export(M(), (torch.rand(4, 4),)).module()
+        module = torch.export.export(M(), (torch.rand(4, 4),), strict=True).module()
 
         tracked_modules: Dict[str, torch.nn.Module] = {
             "module": module,