Fix dmlc#3437 (dmlc#3440)

skorani · Oct 26, 2021 · a9c83bc · a9c83bc
1 parent 579cd3e
commit a9c83bc
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/dgl/backend/pytorch/sparse.py b/python/dgl/backend/pytorch/sparse.py
@@ -102,15 +102,15 @@ def spmm_cache_Y(binary_op, reduce_op, req_grad_X, req_grad_Y):
 
 def spmm_cache_argX(binary_op, reduce_op, req_grad_X, req_grad_Y):
     """Rules to identify whether to cache argX in SpMM forward stage."""
-    if req_grad_X:
+    if req_grad_X or req_grad_Y:
         if reduce_op in ['min', 'max']:
             return True
     return False
 
 
 def spmm_cache_argY(binary_op, reduce_op, req_grad_X, req_grad_Y):
     """Rules to identify whether to cache argY in SpMM forward stage."""
-    if req_grad_Y:
+    if req_grad_X or req_grad_Y:
         if reduce_op in ['min', 'max']:
             return True
     return False