model jutsu

model.py

@@ -82,11 +82,13 @@ class VirTexModel():
             # TODO FIX
             cap_tokens = self.tokenizer.encode(prompt)
             cap_tokens = torch.tensor(cap_tokens, device=self.device).long()
-            subreddit_tokens = subreddit_tokens if sub_prompt is not None else torch.tensor([
-                [self.model.sos_index] +
-                self.tokenizer.encode("pics") +
-                [self.tokenizer.token_to_id("[SEP]")]
-            ])
+            subreddit_tokens = subreddit_tokens if sub_prompt is not None else torch.tensor(
+                (
+                    [self.model.sos_index] + 
+                    self.tokenizer.encode("pics") + 
+                    [self.tokenizer.token_to_id("[SEP]")]
+                ), device = self.device).long()
+
             subreddit_tokens = torch.cat(
                 [
                     subreddit_tokens,
@@ -118,7 +120,7 @@ class VirTexModel():
             else:
                 subreddit, rest_of_caption = "", caption
             
-            is_valid_subreddit = True if sub_prompt is not None or prompt is not None else subreddit in self.valid_subs
+            is_valid_subreddit = subreddit in self.valid_subs
             
             
         return subreddit, rest_of_caption