Update pipeline.py

pipeline.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 from huggingface_hub import PyTorchModelHubMixin
+from PIL import Image
 
 class Generator(nn.Module, PyTorchModelHubMixin):
     def __init__(self, num_channels=3, latent_dim=100, hidden_size=64):
@@ -55,5 +56,5 @@ class PreTrainedPipeline():
             output = self.model(noise)
 
         # Scale image
-        img = transforms.ToPILImage()(output)
+        img = Image.fromarray(output)
         return img