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

use cuda if available

Browse Source
This commit is contained in:
Brett Kuprel
2022-06-28 12:47:11 -04:00
parent 8544f59576
commit 5aa6fe49bf
3 changed files with 13 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
min_dalle/models/dalle_bart_decoder_torch.py
View File

@@ -127,6 +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()
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)
@@ -200,6 +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()
probs, keys_values_state = self.decode_step(
text_tokens = text_tokens,
encoder_state = encoder_state,