updated train.py

saleemhamo · Jul 23, 2024 · 9f95e4a · 9f95e4a
1 parent 5707106
commit 9f95e4a
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/models/fine_grained/train.py b/models/fine_grained/train.py
@@ -113,9 +113,11 @@ def collate_fn(batch):
 
     # Tokenize text sentences
     if config.fine_grained_text_extractor == 'bert':
-        text_sentences = [tokenizer(text, return_tensors='pt', padding='max_length', truncation=True, max_length=512)['input_ids'].squeeze(0) for text in text_sentences]
+        text_sentences = [tokenizer(text, return_tensors='pt', padding='max_length', truncation=True, max_length=512)[
+                              'input_ids'].squeeze(0) for text in text_sentences]
     elif config.fine_grained_text_extractor == 'clip':
-        text_sentences = [tokenizer(texts=[text], return_tensors='pt')['input_ids'].squeeze(0) for text in text_sentences]
+        text_sentences = [tokenizer(text=[text], return_tensors='pt')['input_ids'].squeeze(0) for text in
+                          text_sentences]
 
     # Convert video frames and text sentences to numpy arrays first
     video_frames_padded = pad_sequence([torch.tensor(np.array(v)) for v in video_frames], batch_first=True)