[mypy][5/N] Support all typing on model executor (#4427)

This commit is contained in:
SangBin Cho
2024-04-29 11:01:26 +09:00
committed by GitHub
parent 03dd7d52bf
commit df29793dc7
10 changed files with 61 additions and 34 deletions

View File

@@ -1,5 +1,5 @@
from abc import ABC, abstractmethod
from typing import Any, Dict, List
from typing import Any, Dict, List, Optional
import torch
from torch import nn
@@ -76,8 +76,16 @@ class QuantizationConfig(ABC):
"quantization config.")
@abstractmethod
def get_quant_method(self, layer: torch.nn.Module) -> QuantizeMethodBase:
"""Get the quantize method to use for the quantized layer."""
def get_quant_method(
self, layer: torch.nn.Module) -> Optional[QuantizeMethodBase]:
"""Get the quantize method to use for the quantized layer.
Args:
layer: The layer for the quant method.
Returns:
The quantize method. None if the given layer doesn't support quant
method.
"""
raise NotImplementedError
@abstractmethod