Fixes

app.py

@@ -9,8 +9,8 @@ with gr.Blocks() as demo:
         """
         # ReSWD
 
-        <a href="https://reservoirswd.github.io/"><img src="https://img.shields.io/badge/Project%20Page-5CE1BC.svg"></a> <br>
-        <a href="https://reservoirswd.github.io/static/paper.pdf"><img src="https://img.shields.io/badge/Arxiv-2408.00653-B31B1B.svg"></a> <br>
+        <a href="https://reservoirswd.github.io/"><img src="https://img.shields.io/badge/Project%20Page-5CE1BC.svg"></a>
+        <a href="https://reservoirswd.github.io/static/paper.pdf"><img src="https://img.shields.io/badge/Arxiv-2408.00653-B31B1B.svg"></a>
         <a href="https://github.com/Stability-AI/ReSWD"><img src="https://img.shields.io/badge/git-%23F05032?logo=git&logoColor=white"></a>
         <br>
 

src/color_matcher.py

@@ -85,10 +85,9 @@ def train(
     source_img = source_img.cuda()
 
     batch_size = source_img.shape[0]
-    cdl = CDL(batch_size)
+    cdl = CDL(batch_size).cuda()
 
     optim = torch.optim.Adam(cdl.parameters(), lr=lr)
-    cdl, optim = cdl.cuda(), optim.cuda()
 
     lossses = []
     for i in tqdm(range(num_steps), disable=silent):