AFAIK Python doesn’t have a mechanism for branch prediction because in python and pytorch code is executed eagerly as in line by line
So to make code run faster you can use a jit like torch.jit.script or torch.compile which will specialize on or another branch and there’s some active discussions about explicitly adding control flow ops in torch Add support for dynamic control flow in torch.fx · Issue #99598 · pytorch/pytorch · GitHub