feat: update model

model.py

@@ -1,16 +1,13 @@
-import warnings
 from typing import Optional
+
 import torch
-from torch import nn
-from torch.nn import Module
+from torch import nn, Tensor
 import torch.nn.functional as F
-from torch import Tensor
-from huggingface_hub import PyTorchModelHubMixin
 
-warnings.filterwarnings("ignore", category=UserWarning)
+from huggingface_hub import PyTorchModelHubMixin
 
 
-class Transpose(Module):
+class Transpose(nn.Module):
     def __init__(self, *dims, contiguous=False):
         super(Transpose, self).__init__()
         self.dims, self.contiguous = dims, contiguous