   >(
     _ input: Tensor<T>
   ) -> Tensor<OutType> {
-    return Tensor<OutType>(OutType(input.shape.contiguousSize), on: Device.defaultXLA)
+    return Tensor<OutType>(OutType(input.shape.contiguousSize), on: input.device)
   }
 
   /// Return a slice from 'input'.

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


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

