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

Code before:
            self.logger.info(f"Engine loaded, inputs:{self.engine.input_names}")
        except Exception as e:
            self.logger.info(f"Exception while loading the engine:\n{e}")

    def forward(self, model, argv, kwargs):
        """
        Main forward method:
         Builds TRT engine if not available yet.
         Tries to run TRT engine
         If exception thrown and self.callback==True: falls back to original Pytorch

        Args: Passing through whatever args wrapped module's forward() has
        Returns: Passing through wrapped module's forward() return value(s)

        """
        args = self.defaults
        args.update(kwargs)
        if len(argv) > 0:
            args.update(self._inputs_to_dict(argv))

        if self.engine is None and not self.disabled:
            # Restore original forward for export
            new_forward = model.forward
            model.forward = self.old_forward
            try:
                self._load_engine()
                if self.engine is None:
                    build_args = args.copy()
                    with torch.no_grad():
                        self._build_and_save(model, build_args)
                        # This will reassign input_names from the engine
                    self._load_engine()
                    assert self.engine is not None
            except Exception as e:
                if self.fallback:
                    self.logger.info(f"Failed to build engine: {e}")
                    self.disabled = True
                else:
                    raise e
            if not self.disabled and not self.fallback:
                # Delete all parameters
                for param in model.parameters():
                    del param
                # Call empty_cache to release GPU memory
                torch.cuda.empty_cache()
            # restore TRT hook
            model.forward = new_forward
        # Run the engine
        try:
            if self.engine is not None:
                # forward_trt is not thread safe as we do not use per-thread execution contexts

Review comment (suggestion):
May I ask the reason for adding the `torch.no_grad()` here? Was it caused some issues in the previous version?

Thanks,  
Bin

Code after:

    def forward(self, model, argv, kwargs):
        """
        Main forward method:
         Builds TRT engine if not available yet.
         Tries to run TRT engine
         If exception thrown and self.callback==True: falls back to original Pytorch

        Args: Passing through whatever args wrapped module's forward() has
        Returns: Passing through wrapped module's forward() return value(s)

        """
        args = self.defaults
        args.update(kwargs)
        if len(argv) > 0:
            args.update(self._inputs_to_dict(argv))

        if self.engine is None and not self.disabled:
            # Restore original forward for export
            new_forward = model.forward
            model.forward = self.old_forward
            try:
                self._load_engine()
                if self.engine is None:
                    build_args = args.copy()
                    with torch.no_grad():
                        self._build_and_save(model, build_args)
                        # This will reassign input_names from the engine
                    self._load_engine()
                    assert self.engine is not None
            except Exception as e:
                if self.fallback:
                    self.logger.info(f"Failed to build engine: {e}")
                    self.disabled = True
                else:
                    raise e
            if not self.disabled and not self.fallback:
                # Delete all parameters
                for param in model.parameters():
                    del param
                # Call empty_cache to release GPU memory
                torch.cuda.empty_cache()
            # restore TRT hook
            model.forward = new_forward
        # Run the engine
        try:
            if self.engine is not None:
                # forward_trt is not thread safe as we do not use per-thread execution contexts
                with lock_sm:
                    device = torch.cuda.current_device()
                    stream = torch.cuda.Stream(device=device)