Glenn Jocher
4 years ago
1 changed files with 3 additions and 3 deletions
+ 3
- 3
test.py
View File
| import argparse | import argparse | ||||
| import json | import json | ||||
| import yaml | |||||
| from torch.utils.data import DataLoader | from torch.utils.data import DataLoader | ||||
| from utils.datasets import * | from utils.datasets import * | ||||
| if half: | if half: | ||||
| model.half() # to FP16 | model.half() # to FP16 | ||||
| if device.type != 'cpu' and torch.cuda.device_count() > 1: | |||||
| model = nn.DataParallel(model) | |||||
| # Multi-GPU disabled, incompatible with .half() | |||||
| # if device.type != 'cpu' and torch.cuda.device_count() > 1: | |||||
| # model = nn.DataParallel(model) | |||||
| else: # called by train.py | else: # called by train.py | ||||
| training = True | training = True |
Loading…