def __enter__()

in src/optimum/nvidia/utils/patching.py [0:0]

• 14 lines of code
• 5 McCabe index (conditional complexity)

    def __enter__(self):
        with open(self.builder_path, "r") as file:
            lines = file.readlines()

        for i in range(len(lines)):
            if (
                i > 0
                and lines[i] == "    optimize(network)\n"
                and "# PATCHED" not in lines[i - 1]
            ):
                lines[i] = (
                    "    if getattr(model.config, 'optimize_network', True):\n        # PATCHED\n        optimize(network)\n    else:\n        logger.info('Network optimization disabled during build.')\n"
                )

        with open(self.builder_path, "w") as f:
            f.writelines(lines)