modified: config.json

config.json

@@ -1,9 +1,9 @@
 {
-  "_name_or_path": "asRobotics/fingernet",
-  "architectures": ["FingerNet"],
-  "model_type": "fingernet",
+  "_name_or_path": "asRobotics/ballnet",
+  "architectures": ["BallNet"],
+  "model_type": "ballnet",
   "x_dim": [6],
-  "y_dim": [6, 1800],
+  "y_dim": [6, 2931],
   "h1_dim": [100, 1000],
   "h2_dim": [100, 1000],
   "torch_dtype": "float32",

modeling.py

@@ -2,8 +2,8 @@ import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, PretrainedConfig
 
-class FingerNetConfig(PretrainedConfig):
-    model_type = "fingernet"
+class BallNetConfig(PretrainedConfig):
+    model_type = "ballnet"
 
     def __init__(
         self,
@@ -20,8 +20,8 @@ class FingerNetConfig(PretrainedConfig):
         self.h2_dim = h2_dim
 
 
-class FingerNet(PreTrainedModel):
-    config_class = FingerNetConfig
+class BallNet(PreTrainedModel):
+    config_class = BallNetConfig
 
     def __init__(self, config):
         super().__init__(config)