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

previous commit broke flax model, fixed now

Browse Source
This commit is contained in:
Brett Kuprel
2022-06-28 12:54:58 -04:00
parent 5aa6fe49bf
commit 9d6b6dcc92
4 changed files with 16 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
min_dalle/models/dalle_bart_decoder_torch.py
View File

@@ -127,10 +127,10 @@ class DalleBartDecoderTorch(nn.Module):
self.start_token = torch.tensor([start_token]).to(torch.long)
self.pad_token = torch.tensor([1]).to(torch.long)
self.condition_factor = torch.tensor([10]).to(torch.float)
if torch.cuda.is_available():
self.start_token = self.start_token.cuda()
self.pad_token = self.pad_token.cuda()
self.condition_factor = self.condition_factor.cuda()
# if torch.cuda.is_available():
# self.start_token = self.start_token.cuda()
# self.pad_token = self.pad_token.cuda()
# self.condition_factor = self.condition_factor.cuda()
self.image_token_count = image_token_count
self.embed_tokens = nn.Embedding(image_vocab_size + 1, embed_count)
self.embed_positions = nn.Embedding(image_token_count, embed_count)
@@ -204,7 +204,7 @@ class DalleBartDecoderTorch(nn.Module):
for i in range(self.sample_token_count):
token_index = torch.tensor([i]).to(torch.long)
if torch.cuda.is_available(): token_index = token_index.cuda()
# if torch.cuda.is_available(): token_index = token_index.cuda()
probs, keys_values_state = self.decode_step(
text_tokens = text_tokens,
encoder_state = encoder_state,