def _set_torch_version_environment()

in src/sagemaker_pytorch_container/training.py [0:0]

• 10 lines of code
• 4 McCabe index (conditional complexity)

def _set_torch_version_environment():
    """Set PyTorch version environment variable.

    This is the PyTorch version of the DLC.
    """
    try:
        import torch

        os.environ["SM_DLC_TORCH_VERSION"] = torch.__version__
    except ModuleNotFoundError:
        logger.warn("PyTorch cannot be found")
    except ImportError:
        logger.warn("PyTorch can be found, but cannot be imported")
    except Exception:
        logger.warn("Torch version environment variable cannot be set")