use only 1 gpu if available

JiahuiYu · Mar 8, 2018 · a203e0a · a203e0a · dragonkao730 · Mar 31, 2018
1 parent 8747474
commit a203e0a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test.py b/test.py
@@ -20,6 +20,7 @@
 
 
 if __name__ == "__main__":
+    ng.get_gpus(1)
     args = parser.parse_args()
 
     model = InpaintCAModel()
@@ -29,7 +30,7 @@
     assert image.shape == mask.shape
 
     h, w, _ = image.shape
-    grid = 4
+    grid = 8
     image = image[:h//grid*grid, :w//grid*grid, :]
     mask = mask[:h//grid*grid, :w//grid*grid, :]
     print('Shape of image: {}'.format(image.shape))