Exception is thrown after the job is complete from context_manager_injector.py

SDK version 1.18.0

Submitting a job with TensorFlow estimator and custom docker image. The job completes successfully and the main scripts exits.
But still final status is failed with with this strange error coming from azureml code.

Starting the daemon thread to refresh tokens in background for process with pid = 95

[2020-12-01T23:51:00.945744] The experiment failed. Finalizing run...
12/01/2020 23:51:00 - INFO - __main__ - Exiting context: UserExceptions
12/01/2020 23:51:01 - INFO - __main__ - Exiting context: TrackUserError
12/01/2020 23:51:01 - INFO - __main__ - Exiting context: RunHistory
Cleaning up all outstanding Run operations, waiting 900.0 seconds
2 items cleaning up...
Cleanup took 0.4626293182373047 seconds
12/01/2020 23:51:02 - INFO - __main__ - Exiting context: ProjectPythonPath
Traceback (most recent call last):
File "/mnt/batch/tasks/shared/LS_root/jobs/.../azureml-setup/context_manager_injector.py", line 339, in
execute_with_context(cm_objects, options.invocation)
File "/mnt/batch/tasks/shared/LS_root/jobs/.../azureml-setup/context_manager_injector.py", line 209, in execute_with_context
importlib.reload(imported_module)
File "/opt/miniconda/envs/docker_env/lib/python3.7/importlib/__init__.py", line 148, in reload
raise ImportError(msg.format(name), name=name)
ImportError: module __main__ not in sys.modules

[2020-12-01T23:51:02.097301] Finished context manager injector with Exception.
2020/12/01 23:51:04 Failed to run the wrapper cmd with err: exit status 1
2020/12/01 23:51:04 Attempt 1 of http call to http://10.0.0.4:16384/sendlogstoartifacts/status
2020/12/01 23:51:04 mpirun version string: {
mpirun (Open MPI) 3.1.2

Report bugs to http://www.open-mpi.org/community/help/
}
2020/12/01 23:51:04 Process Exiting with Code: 1