Anton Forsman commited on
Commit
c41ef46
·
1 Parent(s): 414097d
Files changed (1) hide show
  1. diffusion.py +1 -1
diffusion.py CHANGED
@@ -230,4 +230,4 @@ class DiffusionImageAPI:
230
  return [self.get_noisy_image(image, int(t)) for t in time_steps]
231
 
232
  def tensor_to_image(self, tensor):
233
- return Image.fromarray(tensor.cpu().numpy().astype(np.uint8))
 
230
  return [self.get_noisy_image(image, int(t)) for t in time_steps]
231
 
232
  def tensor_to_image(self, tensor):
233
+ return Image.fromarray(tensor.cpu().permute(1, 2, 0).numpy().astype(np.uint8))