Fix some cuda training issues of segmentation (#46)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Laughing
2022-11-17 06:44:02 -06:00
committed by GitHub
parent db1031a1a9
commit 47f1cb3ef4
5 changed files with 38 additions and 21 deletions

View File

@ -6,10 +6,11 @@ from ultralytics.yolo.engine.validator import BaseValidator
class ClassificationValidator(BaseValidator):
def init_metrics(self, model):
self.correct = torch.tensor([])
self.correct = torch.tensor([], device=next(model.parameters()).device)
def preprocess(self, batch):
batch["img"] = batch["img"].to(self.device)
batch["img"] = batch["img"].to(self.device, non_blocking=True)
batch["img"] = batch["img"].half() if self.args.half else batch["img"].float()
batch["cls"] = batch["cls"].to(self.device)
return batch