[test] Removed warning in test modules (#411)

Author: seockho-kim

test/modules/op/bmm.py

@@ -35,7 +35,8 @@ def get_example_inputs(self):
 class SimpleSingleBatchLhsConstBmm(TestModuleBase):
     def __init__(self):
         super().__init__()
-        self.const_lhs = torch.randn(1, 4, 5)
+        const_lhs = torch.randn(1, 4, 5)
+        self.register_buffer("const_lhs", const_lhs)
 
     def forward(self, rhs):
         z = torch.bmm(self.const_lhs, rhs)

test/modules/op/mm.py

@@ -35,7 +35,8 @@ def get_example_inputs(self):
 class SimpleMatmulConstRhs(TestModuleBase):
     def __init__(self):
         super().__init__()
-        self.weight = torch.randn(4, 5)
+        weight = torch.randn(4, 5)
+        self.register_buffer("weight", weight)
 
     def forward(self, lhs):
         out = torch.mm(lhs, self.weight)
@@ -49,7 +50,8 @@ def get_example_inputs(self):
 class SimpleMatmulConstRhsOnert(TestModuleBase):
     def __init__(self):
         super().__init__()
-        self.weight = torch.randn(4, 5)
+        weight = torch.randn(4, 5)
+        self.register_buffer("weight", weight)
 
     def forward(self, lhs):
         out = torch.mm(lhs, self.weight)
@@ -80,7 +82,8 @@ def get_example_inputs(self):
 class SimpleMatmulConstLhsOnertWithLinearConversion(TestModuleBase):
     def __init__(self):
         super().__init__()
-        self.weight = torch.randn(3, 4)
+        weight = torch.randn(3, 4)
+        self.register_buffer("weight", weight)
 
     def forward(self, rhs):
         out = torch.mm(self.weight, rhs)