Unverified Commit 2ebfa6ea authored by xiangyuzhi's avatar xiangyuzhi Committed by GitHub
Browse files

[Sparse]Fix device check (#6409)


Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-24-117.ap-northeast-1.compute.internal>
parent 617979d6
......@@ -61,8 +61,9 @@ void _SDDMMSanityCheck(
mat1.dtype() == mat2.dtype(),
"SDDMM: the two dense matrices should have the same dtype.");
TORCH_CHECK(
mat1.device() == mat2.device(),
"SDDMM: the two dense matrices should on the same device.");
mat1.device() == mat2.device() && sparse_mat->device() == mat2.device(),
"SDDMM: the two dense matrices and sparse matrix should on the same "
"device.");
}
torch::Tensor SDDMMAutoGrad::forward(
......
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