TypeError: 'mode' should be a torch.onnx.TrainingMode enum, but got '<class 'bool'>'.

[jizhang02]

Hello,

I used tensorwatch to test a network,

import tensorwatch as tw
import torchvision.models

alexnet_model = torchvision.models.alexnet()
tw.draw_model(alexnet_model, [1, 3, 224, 224])

but error happened like below:

Is there any solution?