fix wgpu signature

coreylowman · Feb 6, 2024 · 5994ac5 · 5994ac5
1 parent ea424c3
commit 5994ac5
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/dfdx-core/src/tensor_ops/unstack/webgpu_kernel.rs b/dfdx-core/src/tensor_ops/unstack/webgpu_kernel.rs
@@ -1,11 +1,14 @@
 use crate::{prelude::NoneTape, shapes::*, tensor::Webgpu};
 use std::vec::Vec;
 
-impl<E: Dtype, Items> super::UnstackKernel<E, Items> for Webgpu {
-    fn forward<S: Shape>(&self, stack: Tensor<S, E, Self, NoneTape>) -> Result<Items, Error>
+impl<E: Dtype> super::UnstackKernel<E> for Webgpu {
+    fn forward<Items, S: Shape, T: Tape<E, Self>>(
+        &self,
+        stack: Tensor<S, E, Self, NoneTape>,
+    ) -> Result<Items, Error>
     where
         S: super::SubDim,
-        Items: Array<Option<Tensor<S::Tail, E, Self, NoneTape>>, Dim = S::Head>,
+        Items: Array<Tensor<S::Tail, E, Self, T>, Dim = S::Head>,
     {
         todo!()
     }