fix a bug in DCGAN (pytorch#121)

Using single GPU in DCGAN caused an error because 'gpu_ids' in forward function will be None .
ownway22 · Apr 5, 2017 · f89a371 · f89a371
1 parent 0bdcb42
commit f89a371
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dcgan/main.py b/dcgan/main.py
@@ -129,7 +129,7 @@ def __init__(self, ngpu):
 
     def forward(self, input):
         gpu_ids = None
-        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
+        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu >= 1:
             gpu_ids = range(self.ngpu)
         return nn.parallel.data_parallel(self.main, input, gpu_ids)
 
@@ -168,7 +168,7 @@ def __init__(self, ngpu):
 
     def forward(self, input):
         gpu_ids = None
-        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
+        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > =1:
             gpu_ids = range(self.ngpu)
         output = nn.parallel.data_parallel(self.main, input, gpu_ids)
         return output.view(-1, 1)