sagemaker-rapids-exceprt-2.py

from sagemaker.estimator import Estimator

train_instance_type = 'local_gpu'
local_data_dir = 'file://./dataset'

rapids_estimator = Estimator(image_name='sagemaker-rapids:latest',
                          role=role,
                          train_instance_count=1,
                          train_instance_type=train_instance_type,
                          hyperparameters=hyperparams,
                          metric_definitions=[{'Name': 'test_acc', 'Regex': 'test_acc: ([0-9\\.]+)'}])