search

mberr / torch-max-mem

Decorators for maximizing memory utilization with PyTorch & CUDA
https://torch-max-mem.readthedocs.io/en/latest/
MIT License
14 stars 0 forks source link

Code fails for KEYWORD_ONLY params #5

Closed mberr closed 2 years ago

mberr commented 2 years ago

The following snippet

from torch_max_mem import maximize_memory_utilization

@maximize_memory_utilization()
def func(a, *bs, batch_size: int):
    pass

raises an error

Traceback (most recent call last):
  File ".../tmp.py", line 5, in <module>
    def func(a, *bs, batch_size: int):
  File ".../venv/venv-cpu/lib/python3.8/site-packages/torch_max_mem/api.py", line 274, in __call__
    wrapped = maximize_memory_utilization_decorator(
  File ".../venv/venv-cpu/lib/python3.8/site-packages/torch_max_mem/api.py", line 150, in decorator_maximize_memory_utilization
    raise ValueError(f"{parameter_name} must be a keyword based parameter, but is {_parameter.kind}.")
ValueError: batch_size must be a keyword based parameter, but is KEYWORD_ONLY.

since _parameter.kind is KEYWORD_ONLY.

This is overly restrictive, since we only need keyword-based parameters.