Issue with understanding KAN implementation for regression task

[sanjanasur]

I have a regression problem with 2 inputs and 1 output. The input variables have been standardised using minmaxscalar. I tried to follow the HelloKAN program to implement this regression problem but having issues with how to code this problem. Please help with correcting any errors/approach. Size of dataset is 53135.

`X_train, X_test, Y_train, Y_test = train_test_split(dataset, label, test_size=0.25, random_state=76)

Scale the training set

scaler = MinMaxScaler()
X_train = scaler.fit_transform(X_train.to_numpy())

Y_train = Y_train.to_numpy().ravel()

Scale the testing set using the same scaling parameters

X_test = scaler.transform(X_test.to_numpy())

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)from kan import *
torch.set_default_dtype(torch.float64)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

create a KAN: 2D inputs, 1D output, and 5 hidden neurons. cubic spline (k=3), 5 grid intervals (grid=5).

model = KAN(width=[2,5,1], grid=3, k=3, seed=76, device=device)

Assuming X_train and Y_train are already defined and have compatible sizes

Convert X_train and Y_train to PyTorch tensors

X_train_input = torch.tensor(X_train, dtype=torch.float64)
Y_train_label = torch.tensor(Y_train, dtype=torch.float64)

X_test_input = torch.tensor(X_test, dtype=torch.float64)
dataset = {
'train_input': X_train_input,
'train_label': Y_train_label,

'test_input': X_test_input,
'test_label': Y_test[:1000]
}
train the model
model.fit(dataset, opt="LBFGS", steps=50, lamb=0.001);

`
Also I want to evaluate this model using MSE. Please could you guide on how to implement MSE as well for KAN.