Update README.md

README.md

@@ -10,8 +10,6 @@ from lightweight_gan import Generator
 import torch
 from matplotlib import pyplot as plt
 from huggingface_hub import PyTorchModelHubMixin
-class GeneratorWithPyTorchModelHubMixin(gan.__class__, PyTorchModelHubMixin):
-    pass
 
 # Initialize a generator model
 gan_new = Generator(latent_dim=256, image_size=256, attn_res_layers = [32])
@@ -20,6 +18,8 @@ gan_new = Generator(latent_dim=256, image_size=256, attn_res_layers = [32])
 # gan_new.load_state_dict(torch.load('/content/orbgan_e3_state_dict.pt')) 
 
 # Load from model hub:
+class GeneratorWithPyTorchModelHubMixin(gan_new.__class__, PyTorchModelHubMixin):
+    pass
 gan_new.__class__ = GeneratorWithPyTorchModelHubMixin
 gan_new = gan_new.from_pretrained('johnowhitaker/orbgan_e1', latent_dim=256, image_size=256, attn_res_layers = [32])