🩹 Fix MaskedMLP initialization on CUDA device

probabilists · Nov 24, 2024 · b5fc6cc · b5fc6cc
1 parent 2f325e4
commit b5fc6cc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zuko/nn.py b/zuko/nn.py
@@ -270,7 +270,7 @@ def __init__(
         adjacency, inverse = torch.unique(adjacency, dim=0, return_inverse=True)
 
         # P_ij = 1 if A_ik = 1 for all k such that A_jk = 1
-        precedence = adjacency.int() @ adjacency.int().t() == adjacency.sum(dim=-1)
+        precedence = adjacency.double() @ adjacency.double().t() == adjacency.sum(dim=-1)
 
         # Layers
         layers = []