coreylowman · coreylowman · Oct 15, 2022 · Oct 15, 2022
diff --git a/examples/11-conv-net.rs b/examples/11-conv-net.rs
@@ -2,34 +2,31 @@
 //! layers on nightly rust.
 
 #![cfg_attr(feature = "nightly", feature(generic_const_exprs))]
+use dfdx::prelude::*;
+use rand::thread_rng;
 
-#[cfg(not(feature = "nightly"))]
-fn main() {
-    panic!("`+nightly` required to run this example.")
-}
+type Model = (
+    (Conv2D<3, 4, 3>, ReLU),
+    (Conv2D<4, 8, 3>, ReLU),
+    (Conv2D<8, 16, 3>, ReLU),
+    Flatten2D,
+    Linear<7744, 10>,
+);
 
-#[cfg(feature = "nightly")]
 fn main() {
-    use dfdx::prelude::*;
-    use rand::thread_rng;
-
-    type Model = (
-        (Conv2D<3, 4, 3>, ReLU),
-        (Conv2D<4, 8, 3>, ReLU),
-        (Conv2D<8, 16, 3>, ReLU),
-        FlattenImage,
-        Linear<7744, 10>,
-    );
-
     let mut rng = thread_rng();
     let mut m: Model = Default::default();
     m.reset_params(&mut rng);
 
     // single image forward
     let x: Tensor3D<3, 28, 28> = TensorCreator::randn(&mut rng);
+
+    #[cfg(feature = "nightly")]
     let _: Tensor1D<10> = m.forward(x);
 
     // batched image forward
     let x: Tensor4D<32, 3, 28, 28> = TensorCreator::randn(&mut rng);
+
+    #[cfg(feature = "nightly")]
     let _: Tensor2D<32, 10> = m.forward(x);
 }
diff --git a/src/nn/flatten.rs b/src/nn/flatten.rs
@@ -8,25 +8,23 @@ use crate::{Assert, ConstTrue};
 /// Specifically:
 /// ```ignore
 /// # use dfdx::prelude::*;
-/// let _: Tensor1D<{3 * 5 * 7}> = FlattenImage.forward(Tensor3D::<3, 5, 7>::zeros());
-/// let _: Tensor2D<8, {3 * 5 * 7}> = FlattenImage.forward(Tensor4D::<8, 3, 5, 7>::zeros());
+/// let _: Tensor1D<{3 * 5 * 7}> = Flatten2D.forward(Tensor3D::<3, 5, 7>::zeros());
+/// let _: Tensor2D<8, {3 * 5 * 7}> = Flatten2D.forward(Tensor4D::<8, 3, 5, 7>::zeros());
 /// ```
 #[derive(Default, Clone, Copy)]
-pub struct FlattenImage;
+pub struct Flatten2D;
 
-impl ResetParams for FlattenImage {
-    /// Does nothing.
+impl ResetParams for Flatten2D {
     fn reset_params<R: rand::Rng>(&mut self, _: &mut R) {}
 }
 
-impl CanUpdateWithGradients for FlattenImage {
-    /// Does nothing.
+impl CanUpdateWithGradients for Flatten2D {
     fn update<G: GradientProvider>(&mut self, _: &mut G, _: &mut UnusedTensors) {}
 }
 
 #[cfg(feature = "nightly")]
 impl<const M: usize, const N: usize, const O: usize, H: Tape> Module<Tensor3D<M, N, O, H>>
-    for FlattenImage
+    for Flatten2D
 where
     Assert<{ M * N * O == (M * N * O) }>: ConstTrue,
 {
@@ -38,7 +36,7 @@ where
 
 #[cfg(feature = "nightly")]
 impl<const M: usize, const N: usize, const O: usize, const P: usize, H: Tape>
-    Module<Tensor4D<M, N, O, P, H>> for FlattenImage
+    Module<Tensor4D<M, N, O, P, H>> for Flatten2D
 where
     Assert<{ M * N * O * P == M * (N * O * P) }>: ConstTrue,
 {
@@ -48,7 +46,7 @@ where
     }
 }
 
-impl<T> ModuleMut<T> for FlattenImage
+impl<T> ModuleMut<T> for Flatten2D
 where
     Self: Module<T>,
 {
@@ -65,7 +63,7 @@ mod tests {
 
     #[test]
     fn test_flattens() {
-        let _: Tensor1D<{ 15 * 10 * 5 }> = FlattenImage.forward_mut(Tensor3D::<15, 10, 5>::zeros());
-        let _: Tensor2D<5, 24> = FlattenImage.forward_mut(Tensor4D::<5, 4, 3, 2>::zeros());
+        let _: Tensor1D<{ 15 * 10 * 5 }> = Flatten2D.forward_mut(Tensor3D::<15, 10, 5>::zeros());
+        let _: Tensor2D<5, 24> = Flatten2D.forward_mut(Tensor4D::<5, 4, 3, 2>::zeros());
     }
 }
diff --git a/src/nn/npz_impls.rs b/src/nn/npz_impls.rs
@@ -295,7 +295,7 @@ empty_npz_impl!(Dropout);
 empty_npz_impl!(AvgPoolGlobal);
 empty_npz_impl!(MaxPoolGlobal);
 empty_npz_impl!(MinPoolGlobal);
-empty_npz_impl!(FlattenImage);
+empty_npz_impl!(Flatten2D);
 
 impl<const N: usize> SaveToNpz for DropoutOneIn<N> {}
 impl<const N: usize> LoadFromNpz for DropoutOneIn<N> {}