fix fused_mlp import bug

Author: NIRVANALAN

scripts/vit_triplane_sit_sample.py

@@ -11,6 +11,7 @@
 import traceback
 
 import torch as th
+from xformers.components.feedforward import fused_mlp
 # from xformers.triton import FusedLayerNorm as LayerNorm
 import torch.multiprocessing as mp
 import torch.distributed as dist

scripts/vit_triplane_sit_train.py

@@ -34,6 +34,7 @@
 import nsr
 import nsr.lsgm
 # from nsr.train_util_diffusion import TrainLoop3DDiffusion as TrainLoop
+from xformers.components.feedforward import fused_mlp
 
 from datasets.eg3d_dataset import LMDBDataset_MV_Compressed_eg3d
 from nsr.script_util import create_3DAE_model, encoder_and_nsr_defaults, loss_defaults, rendering_options_defaults, eg3d_options_default, dataset_defaults

scripts/vit_triplane_train.py

@@ -15,6 +15,8 @@
 
 import torch as th
 
+from xformers.components.feedforward import fused_mlp
+
 # if th.cuda.is_available(): # FIXME
 #     from xformers.triton import FusedLayerNorm as LayerNorm