deeplearning/fbgemm/fbgemm_gpu/experimental/gen_ai/test/quantize (pytorch#596)

Summary:
Pull Request resolved: facebookresearch/FBGEMM#596

X-link: pytorch#3512

Reviewed By: avikchaudhuri

Differential Revision: D67381311

fbshipit-source-id: 345264f99d6f4b77508b4ea95fe20b3482ad1f04

Author: gmagogsfm

fbgemm_gpu/experimental/gen_ai/test/quantize/quantize_test.py

@@ -134,7 +134,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         model = TestModule().cuda()
         # bf16 required here
-        _ = torch.export.export(model, (torch.randn(32, 32).to(torch.bfloat16).cuda(),))
+        _ = torch.export.export(
+            model, (torch.randn(32, 32).to(torch.bfloat16).cuda(),), strict=True
+        )
 
     def test_f8f8bf16_export(self) -> None:
         class TestModule(torch.nn.Module):
@@ -161,7 +163,7 @@ def forward(self, xq: torch.Tensor, wq: torch.Tensor) -> torch.Tensor:
             fp8_dtype = torch.float8_e4m3fnuz
         xq = torch.randn(M, K).to(fp8_dtype).cuda()
         wq = torch.randn(N, K).to(fp8_dtype).cuda()
-        _ = torch.export.export(model, (xq, wq))
+        _ = torch.export.export(model, (xq, wq), strict=True)
 
 
 @unittest.skipIf(