diff --git a/train_model.py b/train_model.py index ce0d2a5..f08b84e 100644 --- a/train_model.py +++ b/train_model.py @@ -6,6 +6,8 @@ import cv2 from video_compression_model import NUM_CHANNELS, NUM_FRAMES, VideoCompressionModel, PRESET_SPEED_CATEGORIES from tensorflow.keras.callbacks import EarlyStopping +print(tf.config.list_physical_devices('GPU')) + # Constants BATCH_SIZE = 16 EPOCHS = 1