Upload modeling.py with huggingface_hub

modeling.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn as nn
+from typing import List
 from transformers import PreTrainedModel, PretrainedConfig
 
 class BallNetConfig(PretrainedConfig):
@@ -24,16 +25,16 @@ class BallNetConfig(PretrainedConfig):
 
 
 class Normalizer(nn.Module):
-    def __init__(self, mean: Tensor, std: Tensor, eps: float = 1e-8):
+    def __init__(self, mean, std, eps: float = 1e-8):
         super().__init__()
         self.register_buffer("mean", mean)
         self.register_buffer("std", std)
         self.eps = eps
 
-    def normalize(self, x: Tensor) -> Tensor:
+    def normalize(self, x):
         return (x - self.mean) / (self.std + self.eps)
 
-    def denormalize(self, x: Tensor) -> Tensor:
+    def denormalize(self, x):
         return x * (self.std + self.eps) + self.mean
 
 
@@ -108,7 +109,7 @@ class BallNetModel(PreTrainedModel):
 
         self.post_init()
 
-    def forward(self, x: Tensor, **kwargs):
+    def forward(self, x, **kwargs):
         """
         x: (B, 6)
         """