change order of itm loss init

ghstack-source-id: 2e76b2b0df10261b94df452bb2d1f56b4dfffb4f Pull Request resolved: #131
facebookresearch · Jul 23, 2022 · 27c8eec · 27c8eec
1 parent c605279
commit 27c8eec
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/test/models/flava/test_flava.py b/test/models/flava/test_flava.py
@@ -83,7 +83,7 @@ def test_forward_pretraining(self):
  sum(
  value if value is not None else 0 for value in output.losses.values()
  ).item(),
- 21.5150,
+ 21.4029,
  places=4,
  )
 
@@ -107,7 +107,7 @@ def test_forward_pretraining(self):
  sum(
  value if value is not None else 0 for value in output.losses.values()
  ).item(),
- 8.9674,
+ 8.6285,
  places=4,
  )
 
@@ -132,7 +132,7 @@ def test_forward_pretraining(self):
  sum(
  value if value is not None else 0 for value in output.losses.values()
  ).item(),
- 10.0305,
+ 11.0002,
  places=4,
  )
 

diff --git a/torchmultimodal/modules/losses/flava.py b/torchmultimodal/modules/losses/flava.py
@@ -308,7 +308,10 @@ def __init__(
  **kwargs: Any,
  ):
  super().__init__()
-
+ self.itm_loss = ITMLoss(
+ hidden_size=hidden_size,
+ ignore_index=ignore_index,
+ )
  self.contrastive_loss = FLAVAGlobalContrastiveLoss(
  logit_scale=logit_scale,
  image_embedding_size=hidden_size,
@@ -348,10 +351,6 @@ def __init__(
  ),
  }
  )
- self.itm_loss = ITMLoss(
- hidden_size=hidden_size,
- ignore_index=ignore_index,
- )
 
  self.mim_weight = mim_weight
  self.mlm_weight = mlm_weight