Update app.py

app.py

@@ -986,7 +986,7 @@ class BondFloatRBF(torch.nn.Module):
             centers, gamma = self.rbf_params[name]
             rbf = RBF(centers.to(device), gamma.to(device))
             self.rbf_list.append(rbf)
-            linear = torch.nn.Linear(len(centers), embed_dim).cuda()
+            linear = torch.nn.Linear(len(centers), embed_dim).to(device)
             self.linear_list.append(linear)
 
     def forward(self, bond_float_features):