HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Fix some cuda training issues of segmentation (#46)

Browse source 
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
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 38 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

5
ultralytics/yolo/v8/classify/val.py
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