From 5b8d52a6cc4b50d0029cb57f04dddcf65fd7943d Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 8 Dec 2021 23:15:14 +0100 Subject: [PATCH] Default PyTorch Hub to `autocast(False)` (#5926) --- models/common.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/models/common.py b/models/common.py index c269cfef9a6c..b39017378577 100644 --- a/models/common.py +++ b/models/common.py @@ -443,6 +443,7 @@ class AutoShape(nn.Module): multi_label = False # NMS multiple labels per box classes = None # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs max_det = 1000 # maximum number of detections per image + amp = False # Automatic Mixed Precision (AMP) inference def __init__(self, model): super().__init__() @@ -476,8 +477,9 @@ def forward(self, imgs, size=640, augment=False, profile=False): t = [time_sync()] p = next(self.model.parameters()) if self.pt else torch.zeros(1) # for device and type + autocast = self.amp and (p.device.type != 'cpu') # Automatic Mixed Precision (AMP) inference if isinstance(imgs, torch.Tensor): # torch - with amp.autocast(enabled=p.device.type != 'cpu'): + with amp.autocast(enabled=autocast): return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference # Pre-process @@ -506,7 +508,7 @@ def forward(self, imgs, size=640, augment=False, profile=False): x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32 t.append(time_sync()) - with amp.autocast(enabled=p.device.type != 'cpu'): + with amp.autocast(enabled=autocast): # Inference y = self.model(x, augment, profile) # forward t.append(time_sync())