fix python global variable

pytorch/cuda/moe.py

@@ -40,7 +40,7 @@ class MOELayer(nn.Module):
 
     def reset_parameters(self):
         for i in range(self.num_expert):
-            linear = nn.Linear(in_features=self.in_feat, out_features=out_feat)
+            linear = nn.Linear(in_features=self.in_feat, out_features=self.out_feat)
             self.weight.data[i] = linear.weight.data
 
     def forward(self, inp, gate):
@@ -59,7 +59,7 @@ class MOELayer_einsum(nn.Module):
 
     def reset_parameters(self):
         for i in range(self.num_expert):
-            linear = nn.Linear(in_features=self.in_feat, out_features=out_feat)
+            linear = nn.Linear(in_features=self.in_feat, out_features=self.out_feat)
             self.weight.data[i] = linear.weight.data
     
     def forward(self, inp, gate):
@@ -70,24 +70,30 @@ class MOELayer_einsum(nn.Module):
             x[i] = self.weight[gate_long[i]] @ inp[i]
         return x
 
-batch_size = 4
-num_expert = 4
-in_feat = 2
-out_feat = 3
 
-moe = MOELayer(num_expert, in_feat, out_feat).cuda()
-moe_einsum = MOELayer_einsum(num_expert, in_feat, out_feat).cuda()
-moe_einsum.weight.data = moe.weight.data.clone()
+def test():
+    batch_size = 4
+    num_expert = 4
+    in_feat = 2
+    out_feat = 3
 
+    moe = MOELayer(num_expert, in_feat, out_feat).cuda()
+    moe_einsum = MOELayer_einsum(num_expert, in_feat, out_feat).cuda()
+    moe_einsum.weight.data = moe.weight.data.clone()
 
-inp = torch.rand(batch_size, in_feat).cuda()
-gate = torch.randint(low=0, high=num_expert, size=(batch_size, ), requires_grad=False).int().cuda()
 
-output = moe(inp, gate)
-output_einsum = moe_einsum(inp.clone(), gate.clone())
+    inp = torch.rand(batch_size, in_feat).cuda()
+    gate = torch.randint(low=0, high=num_expert, size=(batch_size, ), requires_grad=False).int().cuda()
 
-print(output)
-print(output_einsum)
+    output = moe(inp, gate)
+    output_einsum = moe_einsum(inp.clone(), gate.clone())
 
-#y = output.mean()
-#y.backward()
+    print(output)
+    print(output_einsum)
\ No newline at end of file
+
+    #y = output.mean()
+    #y.backward()
+
+
+if __name__ == '__main__':
+    test()