[FIX][]ValueError: too many values to unpack (expected 5)]

gemini_torch/model.py

@@ -142,20 +142,20 @@ def forward(
             if exists(img):
                 img = self.img_to_transformer(img)
                 x = torch.concat((text, img), dim=1)
-                model_input = self.decoder.forward(x)[0]
+                model_input = self.decoder.forward(x)
 
             # Else, just use the text
             elif exists(img) and exists(audio):
                 # Concat the audio and image and text embeddings all at once, audio is [batch, audio_seq_len]
                 x = self.audio_to_lang_embedding(audio)
                 x = self.img_to_transformer(img)
                 x = torch.concat((text, img, audio), dim=1)
-                model_input = self.decoder.forward(x)[0]
+                model_input = self.decoder.forward(x)
 
             else:
-                model_input = self.decoder.forward(text)[0]
+                model_input = self.decoder.forward(text)
 
-            return self.decoder(model_input, padded_x=model_input[0])
+            return self.decoder(model_input, padded_x=model_input)
         except Exception as e:
             print("Failed in forward method: ", e)
             raise