update seg code
pull/841/head
Kin-Yiu, Wong 2022-09-02 12:21:29 +08:00 committed by GitHub
parent a4df031e45
commit bcc3eff9e2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions

View File

@ -405,7 +405,7 @@ class DetectMultiBackend(nn.Module):
names = model.module.names if hasattr(model, 'module') else model.names # get class names
model.half() if fp16 else model.float()
self.model = model # explicitly assign for to(), cpu(), cuda(), half()
segmentation_model = type(model.model[-1]).__name__ == 'Segment'
segmentation_model = type(model.model[-1]).__name__ in ['Segment', 'ISegment', 'IRSegment']
elif jit: # TorchScript
LOGGER.info(f'Loading {w} for TorchScript inference...')
extra_files = {'config.txt': ''} # model metadata