Skip to content

Commit

Permalink
fixed cuda problem in multimodal_models.py
Browse files Browse the repository at this point in the history
Madjid-CH committed May 17, 2024
1 parent 1ee8fc4 commit e31a234
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion auto_mixer/models/multimodal_models.py
Original file line number Diff line number Diff line change
@@ -143,7 +143,7 @@ def _build_classifiers(self, target_length):

@staticmethod
def _build_encoders(encoders):
return {k: v[1] for k, v in encoders.items()}
return {k: v[1].cuda() for k, v in encoders.items()}

def shared_step(self, batch, **kwargs):
labels = batch['labels']

0 comments on commit e31a234

Please sign in to comment.