About ViT-H, based on your code you were using ViT-H while training, so I understand that you have used it in inference, am I right?
About using ViT-H for inference, it was not a big deal for me, it was just working without any issue, The code for Vit-H looks like that (I used jit trace as it is faster than script TorchScript: Tracing vs. Scripting - Yuxin's Blog)
self.model_scripted = torch.jit.load(model_path).eval().to(device=device_type)
gallery_dataset = SubmissionDataset(
root=self.dataset_path, annotation_file=self.gallery_csv_path,
transforms=get_val_aug_gallery(self.input_size)
)
query_dataset = SubmissionDataset(
root=self.dataset_path, annotation_file=self.queries_csv_path,
transforms=get_val_aug_query(self.input_size), with_bbox=True
)
datasets = ConcatDataset([gallery_dataset, query_dataset])
combine_loader = torch.utils.data.DataLoader(
datasets, batch_size=self.batch_size,
shuffle=False, pin_memory=True, num_workers=self.inference_cfg.num_workers
)
logger.info('Calculating embeddings')
embeddings = []
with torch.cuda.amp.autocast():
with torch.no_grad():
for i, images in tqdm(enumerate(combine_loader), total=len(combine_loader)):
images = images.to(self.device)
outputs = self.model_scripted(images).cpu().numpy()
embeddings.append(outputs)