Update prediction.py

prediction.py

@@ -10,9 +10,9 @@ def predict(prompt, pipe, loss_function=None):
     latents = []
     
     for seed_number, sd_concept in zip(config.SEEDS, config.STABLE_DIFUSION_CONCEPTS):
-        torch.mps.empty_cache()
+        torch.cuda.empty_cache()
         gc.collect()
-        torch.mps.empty_cache()
+        torch.cuda.empty_cache()
         
         prompt = [f'{prompt} {sd_concept}']
         latent = generator.generate_images(pipe=pipe, seed_number=seed_number, prompt=prompt, loss_function=loss_function)