tanner/min-dalle-test
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

delete cache

Browse Source
This commit is contained in:
Brett Kuprel 2022-06-30 15:48:20 -04:00
parent d9d7f34b22
commit 2311a1af7b
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
min_dalle/models/dalle_bart_decoder_flax.py
View File

@ -148,7 +148,7 @@ class DalleBartDecoderFlax(nn.Module):
) )
self.layers = nn.scan( self.layers = nn.scan(
DalleBartDecoderLayerFlax, DalleBartDecoderLayerFlax,
variable_axes = { "params": 0, "cache": 0 }, variable_axes = { "params": 0 },
split_rngs = { "params": True }, split_rngs = { "params": True },
in_axes = (nn.broadcast, 0, nn.broadcast, nn.broadcast), in_axes = (nn.broadcast, 0, nn.broadcast, nn.broadcast),
out_axes = 0, out_axes = 0,

2
min_dalle/models/dalle_bart_encoder_flax.py
View File

@ -124,7 +124,7 @@ class DalleBartEncoderFlax(nn.Module):
self.embed_positions = nn.Embed(self.text_token_count, self.embed_count) self.embed_positions = nn.Embed(self.text_token_count, self.embed_count)
self.layers = nn.scan( self.layers = nn.scan(
DalleBartEncoderLayerFlax, DalleBartEncoderLayerFlax,
variable_axes = { "params": 0, "cache": 0 }, variable_axes = { "params": 0 },
split_rngs = { "params": True }, split_rngs = { "params": True },
in_axes = nn.broadcast, in_axes = nn.broadcast,
length = self.layer_count length = self.layer_count