Dtype fix for backprop matmul (intel#104)

Co-authored-by: SarahByrneIntel <[email protected]>
Nagico2 · Jul 23, 2024 · 86e4e1e · 86e4e1e
1 parent 0562410
commit 86e4e1e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/intel_npu_acceleration_library/nn/autograd.py b/intel_npu_acceleration_library/nn/autograd.py
@@ -63,6 +63,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Iterable[Union[torch.Tensor, Non
 
         dl_dx = run_matmul(grad_output, torch.transpose(w, -1, -2))
         dl_dw = run_matmul(
-            torch.transpose(grad_output, -1, -2), torch.transpose(x, -1, -2)
+            torch.transpose(grad_output, -1, -2),
+            torch.transpose(x, -1, -2).to(torch.float16),
         )
         return dl_dx, dl_dw, None