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

clamp tokens to appropriate bounds

Browse Source
This commit is contained in:
Brett Kuprel
2022-07-08 09:19:28 -04:00
parent e409c120d0
commit 9eb5633931
5 changed files with 13 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
min_dalle/models/dalle_bart_decoder.py
+2
View File
@@ -117,6 +117,7 @@ class DalleBartDecoder(nn.Module):
super().__init__()
self.layer_count = layer_count
self.embed_count = embed_count
self.image_vocab_count = image_vocab_count
self.embed_tokens = nn.Embedding(image_vocab_count + 1, embed_count)
self.embed_positions = nn.Embedding(IMAGE_TOKEN_COUNT, embed_count)
self.layers: List[DecoderLayer] = nn.ModuleList([
@@ -152,6 +153,7 @@ class DalleBartDecoder(nn.Module):
image_count = encoder_state.shape[0] // 2
token_index_batched = token_index[[0] * image_count * 2]
prev_tokens = prev_tokens[list(range(image_count)) * 2]
prev_tokens = prev_tokens.clamp(0, self.image_vocab_count)
decoder_state = self.embed_tokens.forward(prev_tokens)
decoder_state += self.embed_positions.forward(token_index_batched)
decoder_state = self.layernorm_embedding.forward(decoder_state)