What is the `torch.cond` way to write conditionals in PyTorch?

I have learnt, that if you want a graph safe way of using conditionals inside of a PyTorch model you should use torch.cond. It will reduce the painful memcpy (😱) events between the host and device (CUDA) which while algorithmically not a problem, performance wise makes inference very slow. The way I have used torch.cond should be correct (right?) yet also it is a bit “ugly”. I was wondering if there are any best practices that come with using torch.cond to make the model description conc...