Update pipeline.py

pipeline.py

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