torchtnt/tests/framework (#952)

Summary: Pull Request resolved: #952

Reviewed By: galrotem, avikchaudhuri

Differential Revision: D67388469

fbshipit-source-id: 6b0fce0e95c1561407c9e118b2b45d47299c8b56

Author: gmagogsfm

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,