fixing dtype and device in test

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:
pytorch · Oct 22, 2024 · 2b08079 · 2b08079
1 parent 708c698
commit 2b08079
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/integration/test_integration.py b/test/integration/test_integration.py
@@ -1291,8 +1291,8 @@ def forward(self, x):
                 y = self.mha(x, x, x)[0]
                 return self.lin(y)
 
-        mod = MHAModel()
-        input = torch.randn(1,1,4096)
+        mod = MHAModel().to(device).to(dtype)
+        input = torch.randn(1,1,4096).to(device).to(dtype)
         out=mod(*input)
 
         torchao.autoquant(mod, set_inductor_config=False)