".github/vscode:/vscode.git/clone" did not exist on "ca1dc1e7d16958893aa4ef3e005ad419e55a4b71"
Unverified Commit 4d5a96e4 authored by jiqing-feng's avatar jiqing-feng Committed by GitHub
Browse files

fix autocast (#11190)


Signed-off-by: default avatarjiqing-feng <jiqing.feng@intel.com>
parent a7f07c1e
......@@ -221,7 +221,7 @@ class BnB8bitBasicTests(Base8bitTests):
self.assertTrue(module.weight.dtype == torch.int8)
# test if inference works.
with torch.no_grad() and torch.amp.autocast("cuda", dtype=torch.float16):
with torch.no_grad() and torch.autocast(model.device.type, dtype=torch.float16):
input_dict_for_transformer = self.get_dummy_inputs()
model_inputs = {
k: v.to(device=torch_device) for k, v in input_dict_for_transformer.items() if not isinstance(v, bool)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment