Changing nn ToDtype to use generic on method

src/nn/to_dtype.rs

@@ -26,11 +26,12 @@ impl<E1: Dtype, E2: Dtype, D: Device<E1> + Device<E2> + ToDtypeKernel<E1, E2>> T
 }
 
 /// Something that can be copied to have a different dtype
-pub trait ToDtype<E1: Dtype, E2: Dtype, D: Device<E1> + Device<E2> + ToDtypeKernel<E1, E2>>:
-    TensorCollection<E1, D>
-{
-    /// Fallible version of [ToDevice::to_dtype]
-    fn try_to_dtype(&self) -> Result<Self::To<E2, D>, D::Err> {
+pub trait ToDtype<E1: Dtype, D: Device<E1>>: TensorCollection<E1, D> {
+    /// Fallible version of [ToDtype::to_dtype]
+    fn try_to_dtype<E2: Dtype>(&self) -> Result<Self::To<E2, D>, D::Err>
+    where
+        D: Device<E2> + ToDtypeKernel<E1, E2>,
+    {
         let out = Self::iter_tensors(&mut RecursiveWalker {
             m: self,
             f: &mut Converter {
@@ -41,14 +42,45 @@ pub trait ToDtype<E1: Dtype, E2: Dtype, D: Device<E1> + Device<E2> + ToDtypeKern
     }
 
     /// Create a copy of `self` with dtype E2
-    fn to_dtype(&self) -> Self::To<E2, D> {
+    fn to_dtype<E2: Dtype>(&self) -> Self::To<E2, D>
+    where
+        D: Device<E2> + ToDtypeKernel<E1, E2>,
+    {
         self.try_to_dtype().unwrap()
     }
 }
 
-impl<E1: Dtype, E2: Dtype, D: Device<E1> + Device<E2> + ToDtypeKernel<E1, E2>, T> ToDtype<E1, E2, D>
-    for T
-where
-    T: TensorCollection<E1, D>,
-{
+impl<E1: Dtype, D: Device<E1>, T: TensorCollection<E1, D>> ToDtype<E1, D> for T {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{nn::builders::*, tests::*};
+
+    #[test]
+    fn test_linear_f64_to_f32() {
+        let dev: TestDevice = Default::default();
+        type Model = Linear<2, 5>;
+        let model_f64: modules::Linear<2, 5, f64, TestDevice> = dev.build_module::<Model, f64>();
+        let model_f32: modules::Linear<2, 5, f32, TestDevice> = model_f64.to_dtype::<f32>();
+
+        assert_eq!(
+            model_f32.weight.as_vec(),
+            model_f64
+                .weight
+                .as_vec()
+                .iter()
+                .map(|x| *x as f32)
+                .collect::<Vec<_>>()
+        );
+        assert_eq!(
+            model_f32.bias.as_vec(),
+            model_f64
+                .bias
+                .as_vec()
+                .iter()
+                .map(|x| *x as f32)
+                .collect::<Vec<_>>()
+        );
+    }
 }