Update app.py

app.py

@@ -57,6 +57,8 @@ config = OmegaConf.load(CONFIG)
 config = dict2namespace(config)
 
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+config.train.cuda_visible_device = device
 model = NamedCurves(config.model).to(device)
 model.load_state_dict(torch.load(model_pt)["model_state_dict"])