Repository: Project-MONAI/MONAI
File: monai/networks/trt_compiler.py (Python)

Code before:
                cuassert(
                    cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal)
                )
                self.context.execute_async_v3(stream)
                graph = cuassert(cudart.cudaStreamEndCapture(stream))
                self.cuda_graph_instance = cuassert(cudart.cudaGraphInstantiate(graph, 0))
                self.logger.info("CUDA Graph captured!")
        else:
            noerror = self.context.execute_async_v3(stream)
            cuassert(cudart.cudaStreamSynchronize(stream))
            if not noerror:
                raise ValueError("ERROR: inference failed.")

        return self.tensors


def remove_non_tensors(input_example, remove_constants=True):
    #
    # TODO : see if we can instantiate wrappers to handle non-default non-tensors
    #
    non_tensors = {}
    for k, v in input_example.items():
        if v is None:
            non_tensors[k] = v
        elif not torch.is_tensor(v):
            if remove_constants:
                non_tensors[k] = v
            else:
                input_example[k] = torch.tensor(v)

    for key in non_tensors.keys():
        # print(f"Removing non-tensor input: {key})")
        input_example.pop(key)
    return non_tensors


def unroll_input(input_names, input_example):
    # Simulate list/tuple unrolling during ONNX export
    unrolled_input = {}
    for name in input_names:
        val = input_example[name]
        if val is not None:
            if isinstance(val, list | tuple):
                for i in range(len(val)):
                    unrolled_input[f"{name}_{i}"] = val[i]
            else:
                unrolled_input[name] = val
    return unrolled_input


def parse_groups(

Review comment (suggestion):
I think `input_example.get(name, None)` is a better choice here, in case there are any illegal keys.

Thanks,  
Bin

Code after:
                # capture cuda graph
                cuassert(
                    cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal)
                )
                self.context.execute_async_v3(stream)
                graph = cuassert(cudart.cudaStreamEndCapture(stream))
                self.cuda_graph_instance = cuassert(cudart.cudaGraphInstantiate(graph, 0))
                self.logger.info("CUDA Graph captured!")
        else:
            noerror = self.context.execute_async_v3(stream)
            cuassert(cudart.cudaStreamSynchronize(stream))
            if not noerror:
                raise ValueError("ERROR: inference failed.")

        return self.tensors


def make_tensor(d):
    return d if isinstance(d, torch.Tensor) else torch.tensor(d).cuda()


def unroll_input(input_names, input_example):
    # Simulate list/tuple unrolling during ONNX export
    unrolled_input = {}
    for name in input_names:
        val = input_example[name]
        if val is not None:
            if isinstance(val, list) or isinstance(val, tuple):
                for i in range(len(val)):
                    unrolled_input[f"{name}_{i}"] = make_tensor(val[i])
            else:
                unrolled_input[name] = make_tensor(val)
    return unrolled_input


def parse_groups(
    ret: List[torch.Tensor], output_lists: List[List[int]]
) -> Tuple[Union[torch.Tensor, List[torch.Tensor]], ...]:
    """
    Implements parsing of 'output_lists' arg of trt_compile().

    Args:
      ret: plain list of Tensors

      output_lists: list of output group sizes: to form some Lists/Tuples out of 'ret' List, this will be a list
                    of group dimensions, like [[], [5], [-1]] for returning Tensor, list of 5 items and dynamic list.
        Format: [[group_n] | [], ...]
          [] or group_n == 0 : next output from ret is a scalar
          group_n > 0  :       next output from ret is a list of group_n length
          group_n == -1:       next output is a dynamic list. This entry can be at any
                               position in output_lists, but can appear only once.