Decomposing torch.aten.bmm

Hello all,

For bert-base-uncased · Hugging Face model, we got a Torch Fx Graph and called torch_mlir.compile to get a valid torch MLIR. However when we lower it to TOSA, we get an error about not being able to decompose torch.aten.bmm to TOSA ops. Is there a way to decompose torch.aten.bmm in Torch Fx itself so that torch MLIR is free of it to start with? Or else, anyone working on decomposing torch.aten.bmm to Tosa Ops?