Update inference_v5_ensemble.py

2022-04-17 15:39:50 -05:00
parent ae03754942
commit ef58d60482
1 changed files with 8 additions and 0 deletions
--- a/inference_v5_ensemble.py
+++ b/inference_v5_ensemble.py
@@ -501,6 +501,14 @@ def main(window: tk.Wm, text_widget: tk.Text, button_widget: tk.Button, progress
                                _, bin_image = cv2.imencode('.jpg', image)
                                bin_image.tofile(f)
                                
+                        text_widget.write(base_text + 'Clearing CUDA Cache... ') 
+                                
+                        torch.cuda.empty_cache()
+                        time.sleep(3)     
+                        
+                        text_widget.write('Done!\n')
+                                
+                              
                        text_widget.write(base_text + 'Completed Seperation!\n\n')  
                        
                # Emsembling Outputs