Narsil HF staff commited on
Commit
66d0a2c
1 Parent(s): 55fd70d

Fixing scaling ?

Browse files
Files changed (1) hide show
  1. pipeline.py +4 -1
pipeline.py CHANGED
@@ -33,5 +33,8 @@ class PreTrainedPipeline():
33
  output = self.model(noise_vector, class_vector, self.truncation)
34
 
35
  # Ugly hack, I'm sure there's a better way of doing what I want.
36
- img = transforms.ToPILImage()(output[0])
 
 
 
37
  return img
 
33
  output = self.model(noise_vector, class_vector, self.truncation)
34
 
35
  # Ugly hack, I'm sure there's a better way of doing what I want.
36
+ img = output[0]
37
+ img = img.transpose((0, 2, 3, 1))
38
+ img = (img + 1) / 2.0)
39
+ img = transforms.ToPILImage()(img)
40
  return img