Update.

diff --git a/mygpt.py b/mygpt.py
index 90102bf..87071c3 100755 (executable)
--- a/mygpt.py
+++ b/mygpt.py
@@ -16,7 +16,6 @@ import torch, einops
 
 from torch import nn
 from torch.nn import functional as F
-from functorch.dim import dims
 
 import ffutils
 
@@ -352,8 +351,6 @@ class KVRec(nn.Module):
     def forward(self, bs):
         x_q, t0, t1 = bs.x, bs.first, bs.first + bs.nb
 
-        # n,h,l,t,d = dims(5)
-
         if bs.init_cache:
             self.rec_v = x_q.new_zeros(
                 x_q.size(0), self.nb_lines, x_q.size(1), self.w_v.size(1)