From 597f3a01de92be1a4f05114df11d9e39b3608e29 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Sat, 6 Jan 2024 20:06:36 +0100 Subject: [PATCH] Update. --- mygpt.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mygpt.py b/mygpt.py index 4d48247..0e94672 100755 --- a/mygpt.py +++ b/mygpt.py @@ -514,7 +514,7 @@ class Caterpillar(nn.Module): T = bs.x.size(1) DV = self.w_V.size(1) DK = self.w_K.size(1) - Dout = self.w_O.size(1) + DM = self.w_O.size(1) CH = self.caterpillar_height CL = self.caterpillar_length @@ -522,6 +522,8 @@ class Caterpillar(nn.Module): t0 >= CL and (t1 - t0) % CL == 0 ), f"bs.first should be greater than caterpillar_length, and bs.nb should be a multiple of caterpillar_length" + # We cache values to deal efficiently with auto-regression + if bs.init_cache: self.rec_V = X.new_zeros(N, CH, T, DV) self.rec_K = X.new_zeros(N, CH, T, DK) @@ -530,7 +532,7 @@ class Caterpillar(nn.Module): self.rec_V[:, :, t0 - CL : t0] = self.init_V_rec[None, :, :, :] self.rec_K[:, :, t0 - CL : t0] = self.init_K_rec[None, :, :, :] - self.cache_Y = X.new_zeros(N, T, Dout) + self.cache_Y = X.new_zeros(N, T, DM) ###################################################################### # Compute the recurrent state -- 2.39.5