@conda_base(python='3.11.9')

class dummy_test_pypi(FlowSpec):

@batch(gpu=1, image="<http://docker.io/tensorflow/tensorflow:latest-gpu|docker.io/tensorflow/tensorflow:latest-gpu>", queue="job-queue-gpu-metaflow-v2",)

@conda(libraries={'tensorflow': '2.16.1'})

@step

def start(self):

import tensorflow as tf

print("tensorflow" + tf.__version__)

print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

self.next(self.end)

@step

def end(self):

print("All done. \n\n Congratulations!\n")

return