import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
import numpy as np
import os
# Version Information
# tensorflow 2.2.0 , Cudnn7.6.5 and Cuda 10.1 , python 3.8
from keras import backend as K
K.clear_session()

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    tf.config.experimental.set_virtual_device_configuration(gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=6024)])
  except RuntimeError as e:
    print(e)
    os.exit(1)

print(tf.config.experimental.list_physical_devices())

print(tf.__version__)

print(tf.test.is_built_with_cuda())

(X_train, y_train), (X_test,y_test) = tf.keras.datasets.cifar10.load_data()

(X_train, y_train), (X_test,y_test) = tf.keras.datasets.cifar10.load_data()

print(X_train.shape,y_train.shape)