Update train.py

train.py

@@ -66,7 +66,8 @@ def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, sa
         optimizer.step()
 
     torch.save(finetuner.state_dict(), save_path)
-
+    del diffuser
+    torch.cuda.empty_cache()
 if __name__ == '__main__':
 
     import argparse