ml-frameworks/using-mlflow/train-and-deploy-pytorch fails due incompatible torchvision

I'm running https://github.com/Azure/MachineLearningNotebooks/tree/824d844cd7386d95edfa6ecec1642e799ca79dd7/how-to-use-azureml/ml-frameworks/using-mlflow/train-and-deploy-pytorch on a default compute instance with "Python 3.8 AzureML" kernel.

The following cell fails

lib_path = os.path.abspath("scripts")
sys.path.append(lib_path)

import train

with

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
<ipython-input-7-77c0e2068467> in <module>
      2 sys.path.append(lib_path)
      3 
----> 4 import train

/mnt/batch/tasks/shared/LS_root/mounts/clusters/notebook142298/code/how-to-use-azureml/ml-frameworks/using-mlflow/train-and-deploy-pytorch/scripts/train.py in <module>
     15 import torch.nn.functional as F
     16 import torch.optim as optim
---> 17 import torchvision
     18 from torchvision import datasets, transforms
     19 

/anaconda/envs/azureml_py38/lib/python3.8/site-packages/torchvision/__init__.py in <module>
      2 
      3 from torchvision import models
----> 4 from torchvision import datasets
      5 from torchvision import ops
      6 from torchvision import transforms

/anaconda/envs/azureml_py38/lib/python3.8/site-packages/torchvision/datasets/__init__.py in <module>
----> 1 from .lsun import LSUN, LSUNClass
      2 from .folder import ImageFolder, DatasetFolder
      3 from .coco import CocoCaptions, CocoDetection
      4 from .cifar import CIFAR10, CIFAR100
      5 from .stl10 import STL10

/anaconda/envs/azureml_py38/lib/python3.8/site-packages/torchvision/datasets/lsun.py in <module>
     17     import pickle
     18 
---> 19 from .utils import verify_str_arg, iterable_to_str
     20 
     21 

/anaconda/envs/azureml_py38/lib/python3.8/site-packages/torchvision/datasets/utils.py in <module>
      9 import torch
     10 from torch.utils.model_zoo import tqdm
---> 11 from torch._six import PY3
     12 
     13 

ImportError: cannot import name 'PY3' from 'torch._six' (/anaconda/envs/azureml_py38/lib/python3.8/site-packages/torch/_six.py)

Screenshot from 2021-04-09 17-07-33

The conda azureml_py38 environment has the following:

Screenshot from 2021-04-09 17-13-22

Azure / MachineLearningNotebooks

ml-frameworks/using-mlflow/train-and-deploy-pytorch fails due incompatible torchvision #1425