importing spaces

styleTransfer.py

@@ -6,10 +6,12 @@ from PIL import Image  #for importing images
 import torchvision.models as models  #to load vgg 19 model
 import torchvision.transforms as transforms
 from tqdm import tqdm
+import spaces
 
 from dataTransform import load_image
 from vggModel import VGGNet
 
+@spaces.GPU
 def style_transfer(content_img, style_img, total_steps, alpha=1e5, beta=1e10, learning_rate=0.001):
     # Preprocess the input images