Update pipeline.py

pipeline.py

@@ -2,8 +2,9 @@ import torch
 import torch.nn as nn
 from torchvision import transforms
 from PIL import Image
+from huggingface_hub import PyTorchModelHubMixin
 
-class Generator(nn.Module):
+class Generator(nn.Module, PyTorchModelHubMixin):
     def __init__(self):
         super(Generator, self).__init__()
         self.main = nn.Sequential(