def benchmark(self, inputs, count=10, graph_mode=False):
        """
        Benchmark the inference of the AIT model on given inputs, and return the runtime in ms.
        """
        mod = self.module
        if type(inputs) is np.ndarray:
            arr = torch.from_numpy(inputs).cuda()
        else:
            arr = inputs.cuda().contiguous()
        topk = self.cfg.POSTPROCESS.TOPK
        outputs = [
            torch.empty([self.batch_size, 1], dtype=torch.int64).cuda(),
            torch.empty([self.batch_size, topk, 4]).cuda().half(),
            torch.empty([self.batch_size, topk]).cuda().half(),
            torch.empty([self.batch_size, topk], dtype=torch.int64).cuda(),
        ]
        if self.mask_on:
            mask_blob = torch.empty([self.batch_size, topk, 28, 28])
            outputs.append(mask_blob.cuda().half())

        duration, _, _ = mod.benchmark_with_tensors(
            [arr],
            outputs,
            count=count,
            repeat=2,
            graph_mode=graph_mode,
        )
        return duration

        Benchmark the inference of the AIT model on given inputs, and return the runtime in ms.
        
examples/02_detectron2/predictor/predictor.py