Override fixes and general updates (#129)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kalen Michael <kalenmike@gmail.com>
This commit is contained in:
parent
d76d7af566
commit
af6e3c536b
10 changed files with 64 additions and 50 deletions
|
|
@ -107,6 +107,8 @@ class BaseTrainer:
|
|||
self.device = utils.torch_utils.select_device(self.args.device, self.batch_size)
|
||||
self.amp = self.device.type != 'cpu'
|
||||
self.scaler = amp.GradScaler(enabled=self.amp)
|
||||
if self.device.type == 'cpu':
|
||||
self.args.workers = 0 # faster CPU training as time dominated by inference, not dataloading
|
||||
|
||||
# Model and Dataloaders.
|
||||
self.model = self.args.model
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue