def get_torch_version(version):
    """
    returns torch version given a version string. Works for versions like
    "2.1.1", "2.1.1+cpu", "2.1.1+rc" etc and would return 2.1.1 for these
    cases
    """
    version_regex = r"\d+\.\d+\.\d+"
    version = _re.search(version_regex, str(version)).group(0)
    return _StrictVersion(version)

    returns torch version given a version string. Works for versions like
    "2.1.1", "2.1.1+cpu", "2.1.1+rc" etc and would return 2.1.1 for these
    cases
    
coremltools/optimize/torch/_utils/torch_utils.py