diff --git a/server/cros/tradefed/tradefed_test.py b/server/cros/tradefed/tradefed_test.py
index e76e463..f7e45b1 100644
--- a/server/cros/tradefed/tradefed_test.py
+++ b/server/cros/tradefed/tradefed_test.py
@@ -1261,23 +1261,21 @@
                     command = self._tradefed_retry_command(retry_template,
                                                            session_id)
 
-                # TODO(pwang): Evaluate if it is worth it to get the number of
-                #              not-excecuted, for instance, by collecting all
-                #              tests on startup (very expensive, may take 30
-                #              minutes).
                 if media_asset and media_asset.uri:
                     # Clean-up crash logs from previous sessions to ensure
                     # enough disk space for 16GB storage devices: b/156075084.
                     if not keep_media:
                         self._clean_crash_logs()
-                    # TODO(b/137917339): Only prevent screen from turning off for
-                    # media tests. Remove this check once the GPU issue is fixed.
+                # TODO(b/137917339): Only prevent screen from turning off for
+                # media tests. Remove this check once the GPU issue is fixed.
+                keep_screen_on = (media_asset and media_asset.uri) or (
+                        target_module and "Media" in target_module)
+                if keep_screen_on:
                     self._override_powerd_prefs()
                 try:
                     waived_tests, acc = self._run_and_parse_tradefed(command)
                 finally:
-                    # TODO(b/137917339): ditto
-                    if media_asset and media_asset.uri:
+                    if keep_screen_on:
                         self._restore_powerd_prefs()
                 if media_asset:
                     self._fail_on_unexpected_media_download(media_asset)