Commit 479b865b authored by Alessia Marcolini's avatar Alessia Marcolini
Browse files

Add support for multi gpu

parent 25c3d924
......@@ -20,6 +20,14 @@ class ClassificationTileModel(BaseModel):
configuration['n_classes'],
configuration['dropout'],
)
if self.use_cuda:
self.multigpu = True
if self.multigpu:
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
self.resnet = nn.DataParallel(self.resnet)
self.resnet = self.resnet.to(self.device)
if self.is_train: # only defined during training time
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment