diff options
Diffstat (limited to 'candle-core/src/npy.rs')
| -rw-r--r-- | candle-core/src/npy.rs | 8 |
1 files changed, 7 insertions, 1 deletions
diff --git a/candle-core/src/npy.rs b/candle-core/src/npy.rs index 62c62f9e..f3a75965 100644 --- a/candle-core/src/npy.rs +++ b/candle-core/src/npy.rs @@ -85,6 +85,7 @@ impl Header { DType::F16 => "f2", DType::F32 => "f4", DType::F64 => "f8", + DType::I64 => "i8", DType::U32 => "u4", DType::U8 => "u1", }; @@ -160,7 +161,7 @@ impl Header { "f" | "f4" => DType::F32, "d" | "f8" => DType::F64, // "i" | "i4" => DType::S32, - // "q" | "i8" => DType::S64, + "q" | "i8" => DType::I64, // "h" | "i2" => DType::S16, // "b" | "i1" => DType::S8, "B" | "u1" => DType::U8, @@ -233,6 +234,11 @@ impl Tensor { reader.read_u32_into::<LittleEndian>(&mut data_t)?; Tensor::from_vec(data_t, shape, &Device::Cpu) } + DType::I64 => { + let mut data_t = vec![0i64; elem_count]; + reader.read_i64_into::<LittleEndian>(&mut data_t)?; + Tensor::from_vec(data_t, shape, &Device::Cpu) + } } } |