Add additional modules to train a PyTorch model

[rvdinter]

I want to train a PyTorch model. To my understanding, you have to use the following script to train an estimator:

pytorch_estimator = PyTorch('pytorch-train.py',
                            instance_type='ml.p3.2xlarge',
                            instance_count=1,
                            framework_version='1.8.0',
                            py_version='py3',
                            hyperparameters = {'epochs': 20, 'batch-size': 64, 'learning-rate': 0.1})

Here, pytorch_train.py contains all code needed to train the model. However, I am using quite a few custom data augmentation classes, a custom dataset class, training loop, and so on. Is there a possibility to split the file up in multiple modules, and load them in the estimator? If I do this without mentioning in the estimator, I get an error that the modules cannot be found.