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

diff --git a/mygpt.py b/mygpt.py
index 7aa8578..5ea927e 100755
--- a/mygpt.py
+++ b/mygpt.py
@@ -656,21 +656,18 @@ class Caterpillar(nn.Module):
         self.rec_K[:, :, t0:t1] = next_K.flatten(2, 3)
 
         if self.training and self.proba_flashback:
-            # insert_flash_back(
-            # self.rec_V,
-            # V,
-            # self.rec_K,
-            # K,
-            # t0,
-            # t1,
-            # CL,
-            # proba=self.proba_flashback / CL,
-            # )
+            # insert_flash_back(self.rec_V,V,self.rec_K,K,t0,t1,CL,proba=self.proba_flashback / CL,)
+
+            # This piece of code makes the assumption that there is
+            # nothing informative before t0, otherwise we'd have to
+            # implement a cache for V and K too. This should not be
+            # too much of a problem since this is used only during
+            # train, where full sequence are available
 
             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 +676,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