set device correctly

app.py

@@ -37,9 +37,9 @@ def process_image(input_image, learning_rate, iterations, n_targets, seed):
         if pixel_values.grad is not None:
             pixel_values.grad.data.zero_()
 
-        final_activations = get_encoder_activations(pixel_values.to('cuda'))
+        final_activations = get_encoder_activations(pixel_values.to(device))
     
-        logits = model.classifier(final_activations[0]).to(pixel_values.device)
+        logits = model.classifier(final_activations[0]).to(device)
     
         original_loss = BCEWithLogitsLoss(reduction='sum')(logits,random_one_logits)