     _ 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.

    shape: Tensor<T>,
    seed: Tensor<Tseed>
  ) -> Tensor<Dtype> {
    let device = Device.defaultXLA
    return statelessRandomNormal(shape: shape, seed: seed, device: device)
  }


    stop: Tensor<T>,
    num: Tensor<Tidx>
  ) -> Tensor<T> {
    let device = start.device
    return linSpace(start: start, stop: stop, num: num, device: device)
  }

