Buanma

This will be a solution： **device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")** model = KAN(width=[2,5,1], grid=5, k=3, seed=0,**device=device**) dataset = create_dataset(f, n_var=2,**device=device**) model.train(dataset, opt="LBFGS", steps=20, lamb=0.01, lamb_entropy=10.,**device=device**)