Adding graph for a model with MaxUnpool layers

Hello, using torch.utils.tensorboard and tb-nightly I am having some issues adding a graph to the summary. The model I am using is similar to the rgbd ENet implementation (ENetDepth at https://github.com/montrealrobotics/ENetDepth/blob/master/models/enet.py).
I am getting the error:

UserWarning: ONNX export failed on ATen operator max_unpool2d because torch.onnx.symbolic.max_unpool2d does not exist
  .format(op_name, op_name))

The model uses maxunpool layers in the UpsamplingBottlenecks:

self.main_unpool1 = nn.MaxUnpool2d(kernel_size=2)

and forwarded:

        # Main branch shortcut
        main = self.main_conv1(x)
        main = self.main_unpool1(main, max_indices)
        # Extension branch
        ext = self.ext_conv1(x)
        ext = self.ext_conv2(ext)
        ext = self.ext_conv3(ext)
        ext = self.ext_regul(ext)

        # Add main and extension branches
        out = main + ext

        return self.out_prelu(out)

I was wondering if anyone had any tips on a solution to my problem, or a quick workaround?

Thanks