     _ input: Tensor<T>
   ) -> Tensor<OutType> {
     let shape = input.xlaTensor.shape
-    return Tensor(shape.map { OutType($0) }, on: Device.defaultXLA)
+    return Tensor(shape.map { OutType($0) }, on: input.device)
   }
 
   /// Computes sigmoid of `x` element-wise.

    checkSameDevice(x, y)
    checkSamePrecision(x, y)
    let absDiff: Tensor<T> = abs(x - y)
    let dims = Tensor<Int32>(absDiff.shape.dimensions.map { Int32($0) }, on: Device.defaultXLA)
    var value = Tensor<T>(Tensor(tolerance, on: x.device))
    if absDiff.isReducedPrecision {
      value = value.toReducedPrecision

    }
    reduceIdx0.reverse()
    reduceIdx1.reverse()
    let device = s0.device
    return (r0: Tensor(reduceIdx0, on: device), r1: Tensor(reduceIdx1, on: device))
  }

