@conda_base(
    libraries={
        "pytorch::pytorch": "1.13.0",
        "pytorch::torchvision": "0.14.0",
        "conda-forge::boto3": "1.26.107",
        "conda-forge::pandas": "1.5.3",
        "conda-forge::imageio": "2.27.0",
        "conda-forge::torchmetrics": "0.10.2",
        "conda-forge::future": "0.18.2",
        "conda-forge::tensorboard": "2.8.0",
        "conda-forge::pandas-profiling": "3.6.3",
        "conda-forge::bidict": "0.22.1",
        "conda-forge::matplotlib": "3.1",
        "conda-forge::python-graphviz": "0.20.1",
        "conda-forge::torch-scatter": "2.1.1"
    },
    python="3.10.1",
)
class GWFlow(FlowSpec):

    @step
    def start(self):
        """
        This is the start of the flow
        """
        self.next(self.train)

    @enable_decorator(batch(gpu=1, memory=20000), flag=os.getenv("EN_BATCH"))
    @environment(
        vars={
            "EN_BATCH": os.getenv("EN_BATCH"),
            "NVIDIA_DRIVER_CAPABILITIES": "compute,utility",
            "CUDA_VISIBLE_DEVICES": "0,1"
        }
    )
    @step
    def train(self):
         import torch
         import sys
         import os
         from subprocess import call
         # Use cmd to check nvidia card
         print(os.popen("nvidia-smi").read())
         print('__Devices')
         call(["nvidia-smi", "--format=csv",
              "--query-gpu=index,name,driver_version,memory.total,memory.used,memory.free"])
         print(os.popen("nvcc --version").read())

         # See if pytorch picks up the GPUs
         print('__Python VERSION:', sys.version)
         print('__pyTorch VERSION:', torch.__version__)
         print('__CUDA VERSION', torch.version.cuda)
         print('__CUDNN VERSION:', torch.backends.cudnn.version())
         print('__Is CUDA available:', torch.cuda.is_available())
         print('__Number CUDA Devices:', torch.cuda.device_count())
         print('Active CUDA Device: GPU', torch.cuda.current_device())
         print('Available devices ', torch.cuda.device_count())
         print('Current cuda device ', torch.cuda.current_device())
         print(f"GPU count: {torch.cuda.device_count()}")