fix linter errors

import functools

from dataclasses import dataclass

from typing import Callable

from typing_extensions import Literal

import logging

from functools import wraps, cache

from contextlib import contextmanager

import math

import torch.nn as nn