media token fix

open_flamingo/train/train_utils.py

@@ -10,7 +10,17 @@
 from data_utils import DataInfo
 import random
 import numpy as np
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel as DDP
 
+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
 
 def train_one_epoch(
     args,
@@ -77,8 +87,11 @@ def train_one_epoch(
             batch_metadata_to_log[
                 f"{datasets[dataset_ix].name}_num_tokens"
             ] = attention_mask.sum().item()
+            model = unwrap_model(model)
+            model.media_token_id = 400
+            model = DDP(model)
             batch_metadata_to_log[f"{datasets[dataset_ix].name}_num_images"] = (
-                (input_ids == model.media_token_id).sum().item()
+                (input_ids == model.module.media_token_id).sum().item()
             )
 
             # forward pass