Update README.md

README.md

@@ -24,19 +24,20 @@ Find the code in https://github.com/RKorzeniowski/BigBiGAN-PyTorch
 
 ## How to use
 
-def load_model(self,model_path):
-    self.model.load_state_dict(torch.load(model_path))
 
-def init_GAN(self):
-    print('load the pre-trained model')
+    def load_model(self,model_path):
+        self.model.load_state_dict(torch.load(model_path))
 
-    data_set = self.config.ds_name
-    if 'CIFAR10' in data_set:
-        model_path = 'checkpoint_150.pth'
-    elif 'MNIST' in data_set:
-        model_path = 'checkpoint_150.pth'
+    def init_GAN(self):
+        print('load the pre-trained model')
 
-    self.load_model(model_path)
+        data_set = self.config.ds_name
+        if 'CIFAR10' in data_set:
+            model_path = 'checkpoint_150.pth'
+        elif 'MNIST' in data_set:
+            model_path = 'checkpoint_150.pth'
+
+        self.load_model(model_path)