huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-20 12:06:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use a reference for the device.

Browse Source
This commit is contained in:
laurent
2023-06-21 19:55:57 +01:00
parent deb6091099
commit fcb4e6b84f
5 changed files with 23 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
src/tensor.rs
View File

@ -84,7 +84,7 @@ impl Tensor {
fn ones_impl<S: Into<Shape>>(
shape: S,
dtype: DType,
device: Device,
device: &Device,
is_variable: bool,
) -> Result<Self> {
let shape = shape.into();
@ -101,22 +101,22 @@ impl Tensor {
Ok(Self(Arc::new(tensor_)))
}
pub fn ones<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Result<Self> {
pub fn ones<S: Into<Shape>>(shape: S, dtype: DType, device: &Device) -> Result<Self> {
Self::ones_impl(shape, dtype, device, false)
}
pub fn ones_var<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Result<Self> {
pub fn ones_var<S: Into<Shape>>(shape: S, dtype: DType, device: &Device) -> Result<Self> {
Self::ones_impl(shape, dtype, device, true)
}
pub fn ones_like(&self) -> Result<Self> {
Tensor::ones(self.shape(), self.dtype(), self.device())
Tensor::ones(self.shape(), self.dtype(), &self.device())
}
fn zeros_impl<S: Into<Shape>>(
shape: S,
dtype: DType,
device: Device,
device: &Device,
is_variable: bool,
) -> Result<Self> {
let shape = shape.into();
@ -133,21 +133,21 @@ impl Tensor {
Ok(Self(Arc::new(tensor_)))
}
pub fn zeros<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Result<Self> {
pub fn zeros<S: Into<Shape>>(shape: S, dtype: DType, device: &Device) -> Result<Self> {
Self::zeros_impl(shape, dtype, device, false)
}
pub fn zeros_var<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Result<Self> {
pub fn zeros_var<S: Into<Shape>>(shape: S, dtype: DType, device: &Device) -> Result<Self> {
Self::zeros_impl(shape, dtype, device, true)
}
pub fn zeros_like(&self) -> Result<Self> {
Tensor::zeros(self.shape(), self.dtype(), self.device())
Tensor::zeros(self.shape(), self.dtype(), &self.device())
}
pub fn new_impl<A: crate::device::NdArray>(
array: A,
device: Device,
device: &Device,
is_variable: bool,
) -> Result<Self> {
let shape = array.shape()?;
@ -164,11 +164,11 @@ impl Tensor {
Ok(Self(Arc::new(tensor_)))
}
pub fn new<A: crate::device::NdArray>(array: A, device: Device) -> Result<Self> {
pub fn new<A: crate::device::NdArray>(array: A, device: &Device) -> Result<Self> {
Self::new_impl(array, device, false)
}
pub fn var<A: crate::device::NdArray>(array: A, device: Device) -> Result<Self> {
pub fn var<A: crate::device::NdArray>(array: A, device: &Device) -> Result<Self> {
Self::new_impl(array, device, true)
}