apoorva/greenhouse
2
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update utils.py

Browse Source 
solve the multi-gpu training problem
This commit is contained in:
perry0418 2019-03-25 14:59:02 +08:00 committed by GitHub
parent 386835d7ca
commit 4884508110
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
utils/utils.py
View File

@ -285,7 +285,7 @@ def compute_loss(p, targets): # predictions, targets
def build_targets(model, targets, pred): def build_targets(model, targets, pred):
# targets = [image, class, x, y, w, h] # targets = [image, class, x, y, w, h]
if isinstance(model, nn.DataParallel): if isinstance(model, nn.parallel.DistributedDataParallel):
model = model.module model = model.module
yolo_layers = get_yolo_layers(model) yolo_layers = get_yolo_layers(model)