add type to jacobians_per_rank to prevent type missmatch in cpu

arnauqb · Jul 24, 2023 · 44fc911 · 44fc911
1 parent 7e5cf63
commit 44fc911
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/blackbirds/infer/vi.py b/blackbirds/infer/vi.py
@@ -171,7 +171,7 @@ def loss_aux(params):
         if torch.isnan(loss_i) or torch.isnan(jacobian).any():
             continue
         loss += loss_i
-        jacobians_per_rank.append(torch.tensor(jacobian.cpu().numpy()))
+        jacobians_per_rank.append(torch.tensor(jacobian.cpu().numpy(), dtype=torch.float))
         indices_per_rank.append(i)
     # gather the jacobians and parameters from all ranks
     if mpi_size > 1: