Update README.md

README.md

@@ -79,8 +79,8 @@ pretrained_state_dict = torch.load(instantir_path)
 pipe.aggregator.load_state_dict(pretrained_state_dict)
 
 # send to GPU and fp16
-pipe.to(dtype=torch.float16)
-pipe.to('cuda')
+pipe.to(device='cuda', dtype=torch.float16)
+pipe.aggregator.to(device='cuda', dtype=torch.float16)
 ```
 
 Then, you can restore your broken images with: