Right way of comparing tensor sizes without tracer warnings?

In some of my forward functions I have assert statements to check that the tensors passed as input have the expected shape. The asserts work as expected, but when exporting the model to ONNX the tracer throws the warning:

TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert x.shape == y.shape

I have tried wrapping the assert into a with torch.nograd() block and also doing assert list(x.shape) == list(y.shape), but the warning perseveres.

Is there a more standardised approach within pytorch to do these checks?