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

    var broadcastedCondition = condition
    while broadcastedCondition.rank < t.rank {
      broadcastedCondition = expandDims(
        broadcastedCondition, dim: Tensor(Int64(broadcastedCondition.rank), on: Device.defaultXLA))
    }
    broadcastedCondition = broadcastTo(
      broadcastedCondition,
      shape: Tensor<Int32>(t.shape.dimensions.map { Int32($0) }, on: Device.defaultXLA))
    return Tensor(_xla: XLATensor.where_(broadcastedCondition.xlaTensor, t.xlaTensor, e.xlaTensor))
  }


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

