Pip debug fixes (#139)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-01-03 18:54:44 +01:00 committed by GitHub
parent 340376f7a6
commit d74de2582c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 37 additions and 37 deletions

View file

@ -156,7 +156,7 @@ class Exporter:
jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags # export booleans
# Load PyTorch model
self.device = select_device(self.args.device)
self.device = select_device(self.args.device or 'cpu')
if self.args.half:
if self.device.type == 'cpu' or not coreml:
LOGGER.info('half=True only compatible with GPU or CoreML export, i.e. use device=0 or format=coreml')
@ -172,7 +172,9 @@ class Exporter:
# Input
im = torch.zeros(self.args.batch_size, 3, *self.imgsz).to(self.device)
file = Path(getattr(model, 'yaml_file', None) or Path(model.yaml['yaml_file']).name)
file = Path(getattr(model, 'pt_path', None) or model.yaml['yaml_file'])
if file.suffix == '.yaml':
file = Path(file.name)
# Update model
model = deepcopy(model)