def automodel_kwargs()

in optimum_benchmark/backends/pytorch/backend.py [0:0]

• 18 lines of code
• 8 McCabe index (conditional complexity)

    def automodel_kwargs(self) -> Dict[str, Any]:
        kwargs = {}

        if self.config.torch_dtype is not None:
            if hasattr(torch, self.config.torch_dtype):
                kwargs["torch_dtype"] = getattr(torch, self.config.torch_dtype)
            else:
                kwargs["torch_dtype"] = self.config.torch_dtype

        if self.is_quantized:
            kwargs["quantization_config"] = self.quantization_config

        if self.config.attn_implementation is not None:
            kwargs["attn_implementation"] = self.config.attn_implementation

        if self.config.low_cpu_mem_usage is not None:
            kwargs["low_cpu_mem_usage"] = self.config.low_cpu_mem_usage

        if self.config.device_map is not None:
            kwargs["device_map"] = self.config.device_map

        if self.config.tp_plan is not None:
            kwargs["tp_plan"] = self.config.tp_plan

        return kwargs