Migrate from `max.tensor.Tensor` to `max.driver.Tensor`

It looks like `max.tensor` is being deprecated. Mimage makes use of this to store images, so it would make sense to migrate to `max.driver.Tensor` as a return type. 