dnrops
/
burn
mirror of https://github.com/tracel-ai/burn.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

wip

This commit is contained in:
mepatrick73 2024-08-23 16:50:54 -04:00
parent bde8e60dc3
commit 27c7a33dee
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
crates/burn-import/src/burn/node/scatter.rs
View File

@ -10,7 +10,7 @@ pub struct ScatterNode {
pub indices: TensorType, pub indices: TensorType,
pub updates: TensorType, pub updates: TensorType,
pub output: TensorType, pub output: TensorType,
pub axis: usize, axis: usize,
} }
impl<PS: PrecisionSettings> NodeCodegen<PS> for ScatterNode { impl<PS: PrecisionSettings> NodeCodegen<PS> for ScatterNode {
@ -22,6 +22,7 @@ impl<PS: PrecisionSettings> NodeCodegen<PS> for ScatterNode {
vec![ vec![
Type::Tensor(self.input.clone()), Type::Tensor(self.input.clone()),
Type::Tensor(self.indices.clone()), Type::Tensor(self.indices.clone()),
Type::Tensor(self.updates.clone()),
] ]
} }