import os
import sys
import time
from collections import OrderedDict
from time import strftime, gmtime
from tensorboardX import SummaryWriter
from dataset import AsrDataset, DataLoader, AsrCollator
from models.transformer import Model
import hparams as hp
import argparse
import megengine as mge
import megengine.module as M
import megengine.functional as F
from megengine.functional import clip, concat, minimum, norm
from megengine.core._imperative_rt.core2 import pop_scope, push_scope
from typing import Iterable, Union
from megengine.tensor import Tensor
import megengine.distributed as dist
from megengine.data import SequentialSampler, RandomSampler, DataLoader
from criterions.label_smoothing_loss import LabelSmoothingLoss
from megengine.utils.network import Network as Net
import megengine.autodiff as autodiff
import megengine.data as data
import megengine
import multiprocessing

logging = megengine.logger.get_logger()


def clip_grad_norm(
    tensors: Union[Tensor, Iterable[Tensor]],
    max_norm: float,
    ord: float = 2.0,
):
    push_scope("clip_grad_norm")
    if isinstance(tensors, Tensor):
        tensors = [tensors]
    tensors = [t for t in tensors if t.grad is not None]
    norm_ = [norm(t.grad.flatten(), ord=ord) for t in tensors]
    if len(norm_) > 1:
        norm_ = norm(concat(norm_), ord=ord)
    else:
        norm_ = norm_[0]
    scale = max_norm / (norm_ + 1e-6)
    scale = 
minimum(scale, 1)