up

swin2_mose/model.py

@@ -13,8 +13,8 @@ import torch.nn.functional as F
 import torch.utils.checkpoint as checkpoint
 from timm.models.layers import DropPath, to_2tuple, trunc_normal_
 
-from libs import window_reverse, Mlp, window_partition
-from moe import MoE
+from swin2_mose.libs import window_reverse, Mlp, window_partition
+from swin2_mose.moe import MoE
 
 
 class WindowAttention(nn.Module):

swin2_mose/moe.py

@@ -18,7 +18,7 @@ from torch.distributions.normal import Normal
 from copy import deepcopy
 import numpy as np
 
-from libs import Mlp as MLP
+from swin2_mose.libs import Mlp as MLP
 
 
 class SparseDispatcher(object):