diff --git a/main.c b/main.c
index 5a30f4bb26..328c3bbcba 100755
--- a/main.c
+++ b/main.c
@@ -57,6 +57,9 @@
 #include "supervisor/shared/status_leds.h"
 #include "supervisor/shared/stack.h"
 #include "supervisor/serial.h"
+#include "supervisor/usb.h"
+
+#include "shared-bindings/microcontroller/__init__.h"
 
 #include "boards/board.h"
 
@@ -300,6 +303,19 @@ bool run_code_py(safe_mode_t safe_mode) {
         }
     }
 
+    for (uint8_t i = 0; i<=100; i++) {
+        if (!usb_msc_ejected()) {
+            //Go into light sleep
+            break;
+        }
+        mp_hal_delay_ms(10);
+    }
+
+    if (usb_msc_ejected()) {
+        //Go into deep sleep
+        common_hal_mcu_deep_sleep();
+    }
+
     // Display a different completion message if the user has no USB attached (cannot save files)
     if (!serial_connected_at_start) {
         serial_write_compressed(translate("\nCode done running. Waiting for reload.\n"));