import from pytorch dir

pytorch/cuda/moe.py

@@ -2,7 +2,7 @@ import math
 from torch import nn
 import torch
 
-from moe_function import moe
+from .moe_function import moe
 
 
 class MOELayer(nn.Module):

pytorch/mem_transformer.py

@@ -34,7 +34,7 @@ class PositionalEmbedding(nn.Module):
             return pos_emb[:,None,:]
 
 class CustomizedMoEPositionwiseFF(nn.Module):
-    def __init__(self, d_model, d_inner, dropout, pre_lnorm=False, top_k=2, num_expert=32):
+    def __init__(self, d_model, d_inner, dropout, pre_lnorm=False, top_k=2, num_expert=64):
         super(CustomizedMoEPositionwiseFF, self).__init__()
         print("CustomizedMoEPositionwiseFF num_expert=%d top_k=%d" % (num_expert, top_k))