Set workers=0 for MPS Train and Val modes (#4697)

This commit is contained in:
Glenn Jocher 2023-09-02 04:25:03 +02:00 committed by GitHub
parent 2bc6e647c7
commit a4fabfdacf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 88 additions and 74 deletions

View file

@ -107,7 +107,7 @@ class BaseTrainer:
print_args(vars(self.args))
# Device
if self.device.type == 'cpu':
if self.device.type in ('cpu', 'mps'):
self.args.workers = 0 # faster CPU training as time dominated by inference, not dataloading
# Model and Dataset

View file

@ -144,7 +144,7 @@ class BaseValidator:
else:
raise FileNotFoundError(emojis(f"Dataset '{self.args.data}' for task={self.args.task} not found ❌"))
if self.device.type == 'cpu':
if self.device.type in ('cpu', 'mps'):
self.args.workers = 0 # faster CPU val as time dominated by inference, not dataloading
if not pt:
self.args.rect = False