fix: cast torch model output dtype to match input

fix: cast torch model output dtype to match input #2889