diff options
Diffstat (limited to 'util/image_pool.py')
| -rw-r--r-- | util/image_pool.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/util/image_pool.py b/util/image_pool.py index 152ef5b..9f34a09 100644 --- a/util/image_pool.py +++ b/util/image_pool.py @@ -13,7 +13,7 @@ class ImagePool(): if self.pool_size == 0: return images return_images = [] - for image in images.data: + for image in images: image = torch.unsqueeze(image, 0) if self.num_imgs < self.pool_size: self.num_imgs = self.num_imgs + 1 |