Move model class outside of main

forward_UpDimV2_long.py

@@ -17,14 +17,7 @@ from tqdm import tqdm, trange
 from math import ceil
 
 
-def main(args):
-    batch_size = 64
-    num_feature = 4096
-    num_classes = 10
-    rng = np.random.RandomState(42)
-
 class UpDimV2(torch.nn.Module):
-
     def __init__(self, num_class):
         super(UpDimV2, self).__init__()
         self.activation = torch.nn.LeakyReLU(0.001, inplace=True)
@@ -142,6 +135,12 @@ def main(args):
         return self.fc3(out)
 
 
+def main(args):
+    batch_size = 64
+    num_feature = 4096
+    num_classes = 10
+    rng = np.random.RandomState(42)
+
     model = torch.nn.DataParallel(UpDimV2(num_classes))
     model.load_state_dict((torch.load(args.weight)['model']))
     model.to('cuda')