fsoest:

return A_real @ B_real - A_imag @ B_imag + 1j * (A_imag @ B_real + A_real @ B_imag)


I had the same issue with custom functions using float tensors to return complex ones, I solved it by explicitly casting to complex the real and imaginary parts. You can try:
return (A_real @ B_real - A_imag @ B_imag).type(torch.complex64)  \
          + 1j * (A_imag @ B_real + A_real @ B_imag).type(torch.complex64)