File size: 405 Bytes
3461c19 | 1 2 3 4 5 6 7 8 9 10 11 12 13 | # SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
class LightweightModule:
"""Torch modules add a surprising amount of host overhead for attribute
access and method calls. This class is a lightweight alternative that
just wraps a forward function for now."""
def __call__(self, *args, **kwargs):
return self.forward(*args, **kwargs)
|