Merge pull request rapidsai#4263 from rapidsai/branch-24.04

Forward-merge branch-24.04 to branch-24.06
alexbarghi-nv · Mar 20, 2024 · 9d7b5be · 9d7b5be
2 parents b2a3890 + 4becbe8
commit 9d7b5be
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py b/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py
@@ -251,7 +251,10 @@ def forward(
         if edge_envelope is not None:
             out = out * edge_envelope.view(-1, 1)
 
-        out = scatter_reduce(out, dst, dim=0, dim_size=num_dst_nodes, reduce=reduce)
+        dtype = out.dtype
+        out = scatter_reduce(
+            out.float(), dst, dim=0, dim_size=num_dst_nodes, reduce=reduce
+        ).to(dtype)
 
         if self.batch_norm:
             out = self.batch_norm(out)