fix: no need to dtype A in jamba (#32924)

Commit

1 year ago

fix: no need to dtype A in jamba (#32924) Co-authored-by: Gal Cohen <galc@ai21.com>