delete cache

2022-06-30 15:48:20 -04:00
parent d9d7f34b22
commit 2311a1af7b
2 changed files with 2 additions and 2 deletions
@@ -148,7 +148,7 @@ class DalleBartDecoderFlax(nn.Module):
        )
        self.layers = nn.scan(
            DalleBartDecoderLayerFlax,
-            variable_axes = { "params": 0, "cache": 0 },
+            variable_axes = { "params": 0 },
            split_rngs = { "params": True },
            in_axes = (nn.broadcast, 0, nn.broadcast, nn.broadcast),
            out_axes = 0,