bugfix

choderalab · Aug 18, 2024 · f743d74 · f743d74
1 parent 897d1b9
commit f743d74
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/modelforge/potential/processing.py b/modelforge/potential/processing.py
@@ -349,7 +349,9 @@ def physnet_charge_conservation(
 
         # for each atom i, calculate the sum of partial charges for all other
         predicted_per_molecule_charge = torch.zeros(
-            total_charges.shape, dtype=torch.float32, device=total_charges.device
+            total_charges.shape,
+            dtype=per_atom_charge.dtype,
+            device=total_charges.device,
         ).scatter_add_(0, mol_indices.long(), per_atom_charge)
 
         # Calculate the correction factor for each molecule