Merge pull request #283 from mlfoundations/media_token_fix

train_utils media token fix

open_flamingo/train/train_utils.py

@@ -10,6 +10,7 @@
 from data_utils import DataInfo
 import random
 import numpy as np
+import torch.nn as nn
 
 
 def train_one_epoch(
@@ -78,7 +79,7 @@ def train_one_epoch(
                 f"{datasets[dataset_ix].name}_num_tokens"
             ] = attention_mask.sum().item()
             batch_metadata_to_log[f"{datasets[dataset_ix].name}_num_images"] = (
-                (input_ids == model.media_token_id).sum().item()
+                (input_ids == unwrap_model(model).media_token_id).sum().item()
             )
 
             # forward pass
@@ -188,6 +189,16 @@ def random_seed(seed=42, rank=0):
     random.seed(seed + rank)
 
 
+def unwrap_model(model):
+    """
+    Unwrap a model from a DataParallel or DistributedDataParallel wrapper.
+    """
+    if isinstance(model, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
+        return model.module
+    else:
+        return model
+
+
 ################################
 # Helper functions for logging #
 ################################