ColossalAI/colossalai/utils/multi_tensor_apply/__init__.py

4 lines
101 B
Python

from .multi_tensor_apply import MultiTensorApply
multi_tensor_applier = MultiTensorApply(2048 * 32)