ai_dataloader/collate/torch_collate/
nonzero.rs

1use super::super::Collate;
2use super::TorchCollate;
3use ndarray::{Array, Ix1};
4use std::num::{
5    NonZeroI128, NonZeroI16, NonZeroI32, NonZeroI64, NonZeroI8, NonZeroIsize, NonZeroU128,
6    NonZeroU16, NonZeroU32, NonZeroU64, NonZeroU8, NonZeroUsize,
7};
8
9/// Maybe this one may not be supported by a tensor running on a GPU.
10
11macro_rules! nonzero_impl {
12    ($($t:ty)*) => {
13        $(
14            impl Collate<$t> for TorchCollate {
15                type Output = Array<$t, Ix1>;
16                fn collate(&self, batch: Vec<$t>) -> Self::Output {
17                    Array::from_vec(batch)
18                }
19            }
20        )*
21    };
22}
23nonzero_impl!(
24    NonZeroUsize NonZeroU8 NonZeroU16 NonZeroU32 NonZeroU64 NonZeroU128
25    NonZeroIsize NonZeroI8 NonZeroI16 NonZeroI32 NonZeroI64 NonZeroI128
26);