diff options
| author | Jun-Yan Zhu <junyanz@users.noreply.github.com> | 2017-11-02 15:29:22 -0700 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2017-11-02 15:29:22 -0700 |
| commit | d8cb73153cc460f9602b802dade2cff5946d0ac7 (patch) | |
| tree | e822bc83fee4c4f47bbabc528a43f871e1abb5a5 /util/image_pool.py | |
| parent | c9f1962642c63772cd182f93e81908c90ba267f9 (diff) | |
| parent | df16bf1b18e03213efe9b22a0155b833a06c6a9c (diff) | |
Merge pull request #141 from jpmerc/patch-2
gpu memory leaks
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 |