Merge branch 'main' into serialize_optimizers

coreylowman · Jul 10, 2023 · cbcc691 · cbcc691
2 parents 1d31c57 + 442b4a6
commit cbcc691
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/nn/conv.rs b/src/nn/conv.rs
@@ -108,7 +108,8 @@ where
                 |s| &s.weight,
                 |s| &mut s.weight,
                 TensorOptions::reset_with(|t| {
-                    let b = E::ONE / E::from_usize(I * K * K).unwrap().sqrt();
+                    let scale = E::from_f64(G as f64 / (I * K * K) as f64).unwrap();
+                    let b = scale.sqrt();
                     t.try_fill_with_distr(rand_distr::Uniform::new(-b, b))
                 }),
             ),

diff --git a/src/nn/convtrans.rs b/src/nn/convtrans.rs
@@ -104,7 +104,8 @@ where
                 |s| &s.weight,
                 |s| &mut s.weight,
                 TensorOptions::reset_with(|t| {
-                    let b = E::ONE / E::from_usize(I * K * K).unwrap().sqrt();
+                    let scale = E::from_f64(G as f64 / (I * K * K) as f64).unwrap();
+                    let b = scale.sqrt();
                     t.try_fill_with_distr(rand_distr::Uniform::new(-b, b))
                 }),
             ),