cast device

725cc729 · justheuristic · 28a9313d · 725cc729
Commit 725cc729 authored Sep 18, 2022 by justheuristic
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

tests/test_modules.py tests/test_modules.py +1 -1

No files found.
--- a/tests/test_modules.py
+++ b/tests/test_modules.py
@@ -567,7 +567,7 @@ def test_linear8bitlt_no_fp16_weights(threshold, memory_efficient_backward):

        mlp.zero_grad()
        (o1 * grad_proj).sum().backward()
-        grad_ref = grad_proj.flatten(2) @ w2.to() @ w1.to(grad_proj.device)
+        grad_ref = grad_proj.flatten(2) @ w2 @ w1
        assert torch.allclose(b1.grad, grad_ref)