X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=mygpt.py;h=24ba34591bb6833e9f0a550a4e104074486cec4a;hb=43fbfaac1850098f5b1a9470c8e6ca3d5ab479fe;hp=7aa85782e2195849fa2b28803762232b86b79f2f;hpb=359cf44b609cebd0f01b9c2d2be1f76a4577a97b;p=mygptrnn.git

diff --git a/mygpt.py b/mygpt.py
index 7aa8578..24ba345 100755
--- a/mygpt.py
+++ b/mygpt.py
@@ -669,8 +669,8 @@ class Caterpillar(nn.Module):
 
             n = torch.arange(N, device=X.device)[:, None, None, None]
             t = torch.arange(t0, t1, device=X.device)[None, None, :, None]
-            dv = torch.arange(DV)[None, None, None, :]
-            dk = torch.arange(DK)[None, None, None, :]
+            dv = torch.arange(DV, device=X.device)[None, None, None, :]
+            dk = torch.arange(DK, device=X.device)[None, None, None, :]
 
             u = (
                 torch.rand(N, CH, t1 - t0, 1, device=X.device).mul(t).long() // CL
@@ -679,20 +679,22 @@ class Caterpillar(nn.Module):
             src_time = t - u - t0
             src_head = torch.randint(H, (N, CH, t1 - t0, 1), device=X.device)
 
-            mask_V = (torch.rand(N, CH, t1 - t0, DV) <= self.proba_flashback).long()
+            mask_V = (
+                torch.rand(N, CH, t1 - t0, DV, device=X.device) <= self.proba_flashback
+            ).long()
             self.rec_V[:, :, t0:t1] = (
                 mask_V * V[n, src_head, src_time, dv]
                 + (1 - mask_V) * self.rec_V[:, :, t0:t1]
             )
 
-            mask_K = (torch.rand(N, CH, t1 - t0, DK) <= self.proba_flashback).long()
+            mask_K = (
+                torch.rand(N, CH, t1 - t0, DK, device=X.device) <= self.proba_flashback
+            ).long()
             self.rec_K[:, :, t0:t1] = (
                 mask_K * K[n, src_head, src_time, dk]
                 + (1 - mask_K) * self.rec_K[:, :, t0:t1]
             )
 
-        exit(0)
-
         ######################################################################
         # compute the readout