it does not work

Author: lucidrains

README.md

@@ -142,12 +142,3 @@ trainer()
     url     = {https://api.semanticscholar.org/CorpusID:270878436}
 }
 ```
-
-```bibtex
-@inproceedings{Yao2024FasterDiTTF,
-    title   = {FasterDiT: Towards Faster Diffusion Transformers Training without Architecture Modification},
-    author  = {Jingfeng Yao and Wang Cheng and Wenyu Liu and Xinggang Wang},
-    year    = {2024},
-    url     = {https://api.semanticscholar.org/CorpusID:273346237}
-}
-```

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rectified-flow-pytorch"
-version = "0.1.11"
+version = "0.1.12"
 description = "Rectified Flow in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

rectified_flow_pytorch/rectified_flow.py

@@ -116,23 +116,6 @@ class MSELoss(Module):
     def forward(self, pred, target, **kwargs):
         return F.mse_loss(pred, target)
 
-class MSEAndDirectionLoss(Module):
-    """
-    Figure 7 - https://arxiv.org/abs/2410.10356
-    """
-
-    def __init__(self, cosine_sim_dim: int = 1):
-        super().__init__()
-        assert cosine_sim_dim > 0, 'cannot be batch dimension'
-        self.cosine_sim_dim = cosine_sim_dim
-
-    def forward(self, pred, target, **kwargs):
-        mse_loss = F.mse_loss(pred, target)
-
-        direction_loss = (1. - F.cosine_similarity(pred, target, dim = self.cosine_sim_dim)).mean()
-
-        return mse_loss + direction_loss
-
 # loss breakdown
 
 LossBreakdown = namedtuple('LossBreakdown', ['total', 'main', 'data_match', 'velocity_match'])
@@ -152,7 +135,6 @@ def __init__(
         predict: Literal['flow', 'noise'] = 'flow',
         loss_fn: Literal[
             'mse',
-            'mse_and_direction',
             'pseudo_huber',
             'pseudo_huber_with_lpips'
         ] | Module = 'mse',
@@ -197,9 +179,6 @@ def __init__(
         if loss_fn == 'mse':
             loss_fn = MSELoss()
 
-        elif loss_fn == 'mse_and_direction':
-            loss_fn = MSEAndDirectionLoss(**loss_fn_kwargs)
-
         elif loss_fn == 'pseudo_huber':
             assert predict == 'flow'