ClearML/clearml
1
0
Fork 0
mirror of https://github.com/clearml/clearml synced 2025-03-03 10:42:00 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Remove cast warning in Tensorflow v2 example

Browse Source
This commit is contained in:
allegroai 2020-04-28 10:45:02 +03:00
parent 3d9683f290
commit 824808d38b
1 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
examples/tensorflow_v2_mnist.py
View File

@ -19,8 +19,8 @@ mnist = tf.keras.datasets.mnist
x_train, x_test = x_train / 255.0, x_test / 255.0
# Add a channels dimension
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]
x_train = x_train[..., tf.newaxis].astype('float32')
x_test = x_test[..., tf.newaxis].astype('float32')
# Use tf.data to batch and shuffle the dataset
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(32)
@ -31,10 +31,10 @@ test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)
class MyModel(Model):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = Conv2D(32, 3, activation='relu')
self.conv1 = Conv2D(32, 3, activation='relu', dtype=tf.float32)
self.flatten = Flatten()
self.d1 = Dense(128, activation='relu')
self.d2 = Dense(10, activation='softmax')
self.d1 = Dense(128, activation='relu', dtype=tf.float32)
self.d2 = Dense(10, activation='softmax', dtype=tf.float32)
def call(self, x):
x = self.conv1(x)