ugly modify fit_params() in utils.py to use symbolic (in ex3 moon case)

[Simba1999]

add device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') before for loop

and put a_grid and b_grid to device post_fun = fun(a_grid[None,:,:].to(device) * x[:,None,None] + b_grid[None,:,:].to(device))

and before line post_fun = torch.nan_to_num(post_fun), put post_fun and y back to cpu post_fun = post_fun.cpu() y = y.cpu()

may be some smarter way...

[Simba1999]

In moon.py

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

And also put dataset to device

dataset['train_input'] = torch.from_numpy(train_input.astype(np.float32)).to(device) dataset['test_input'] = torch.from_numpy(test_input.astype(np.float32)).to(device) dataset['train_label'] = torch.from_numpy(train_label[:,None]).to(device) dataset['test_label'] = torch.from_numpy(test_label[:,None]).to(device)

modify the plt plt.scatter(X[:,0].cpu(), X[:,1].cpu(), c=y[:,0].cpu())

put the model to device model = KAN(width=[2,1], grid=3, k=3, device=device)

add device after train() results = model.train(dataset, opt="LBFGS", steps=1, metrics=(train_acc, test_acc), device=device)

[Simba1999]

In regression case, in acc(), put y to cpu y = y.cpu().numpy()