Skip to content

Commit

Permalink
Fix two small issues with seanet layers in streaming mode.
Browse files Browse the repository at this point in the history
  • Loading branch information
LaurentMazare committed Dec 2, 2024
1 parent d2d98d0 commit 6a9a3d4
Showing 1 changed file with 3 additions and 6 deletions.
9 changes: 3 additions & 6 deletions rust/moshi-core/src/seanet.rs
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ impl Module for SeaNetResnetBlock {

impl StreamingModule for SeaNetResnetBlock {
fn reset_state(&mut self) {
// TODO(laurent): self.skip_op should probably be resetted here.
self.skip_op.reset_state();
for block in self.block.iter_mut() {
block.reset_state()
}
Expand All @@ -132,12 +132,9 @@ impl StreamingModule for SeaNetResnetBlock {
for block in self.block.iter_mut() {
ys = block.step(&ys.apply(&self.activation)?)?;
}
match self.shortcut.as_ref() {
match self.shortcut.as_mut() {
None => self.skip_op.step(&ys, xs),
Some(shortcut) => {
// TODO(laurent): shouldn't this use shortcut.step(xs) instead?
self.skip_op.step(&ys, &xs.apply(shortcut)?)
}
Some(shortcut) => self.skip_op.step(&ys, &shortcut.step(xs)?),
}
}
}
Expand Down

0 comments on commit 6a9a3d4

Please sign in to comment.