Sophie98 commited on
Commit
a53d88d
1 Parent(s): c10dead

still error

Browse files
Files changed (1) hide show
  1. styleTransfer.py +2 -2
styleTransfer.py CHANGED
@@ -109,9 +109,9 @@ def StyleTransformer(content_img: Image, style_img: Image,
109
  print(type(output))
110
  output = output[0].cpu()
111
  print(type(output))
112
- print(output.shape)
113
  torch2PIL = transforms.ToPILImage()
114
- output = torch2PIL(output)
115
  return output
116
 
117
  ############################################## STYLE-GAN #############################################
 
109
  print(type(output))
110
  output = output[0].cpu()
111
  print(type(output))
112
+ print(output.squeeze().shape)
113
  torch2PIL = transforms.ToPILImage()
114
+ output = torch2PIL(output.squeeze())
115
  return output
116
 
117
  ############################################## STYLE-GAN #############################################