Update.

[mygptrnn.git] / mygpt.py

diff --git a/mygpt.py b/mygpt.py
index 7aa8578..6e13ff8 100755 (executable)
--- a/mygpt.py
+++ b/mygpt.py
@@ -656,6 +656,12 @@ class Caterpillar(nn.Module):
         self.rec_K[:, :, t0:t1] = next_K.flatten(2, 3)
 
         if self.training and self.proba_flashback:
+            # 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
+
             # insert_flash_back(
             # self.rec_V,
             # V,
@@ -669,8 +675,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 +685,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