.

safety_checker_improved.py

@@ -11,9 +11,10 @@ import tensorflow as tf
 from tensorflow.keras import mixed_precision
 physical_devices = tf.config.list_physical_devices('GPU')
 
-tf.config.experimental.set_memory_growth(
-    physical_devices[0], True
-)
+if len(physical_devices) > 0:
+    tf.config.experimental.set_memory_growth(
+        physical_devices[0], True
+    )
 
 model = tf.keras.models.load_model('nsfweffnetv2-b02-3epochs.h5',custom_objects={"KerasLayer":hub.KerasLayer})
 # "The image classifier had been trained on 682550 images from the 5 classes "Drawing" (39026), "Hentai" (28134), "Neutral" (369507), "Porn" (207969) & "Sexy" (37914).