Make func cloneable. (#1137)

This commit is contained in:
Laurent Mazare
2023-10-20 16:28:50 +01:00
committed by GitHub
parent 99cf13e8e2
commit 7366aeac21
3 changed files with 9 additions and 7 deletions

View File

@ -108,7 +108,7 @@ pub fn parse_config<T: AsRef<Path>>(path: T) -> Result<Darknet> {
} }
enum Bl { enum Bl {
Layer(Box<dyn candle_nn::Module + Send>), Layer(Box<dyn candle_nn::Module + Send + Sync>),
Route(Vec<usize>), Route(Vec<usize>),
Shortcut(usize), Shortcut(usize),
Yolo(usize, Vec<(usize, usize)>), Yolo(usize, Vec<(usize, usize)>),

View File

@ -1,10 +1,12 @@
//! Layers defined by closures. //! Layers defined by closures.
use candle::{Result, Tensor}; use candle::{Result, Tensor};
use std::sync::Arc;
/// A layer defined by a simple closure. /// A layer defined by a simple closure.
#[derive(Clone)]
pub struct Func<'a> { pub struct Func<'a> {
#[allow(clippy::type_complexity)] #[allow(clippy::type_complexity)]
f: Box<dyn 'a + Fn(&Tensor) -> Result<Tensor> + Send>, f: Arc<dyn 'a + Fn(&Tensor) -> Result<Tensor> + Send + Sync>,
} }
impl<'a> std::fmt::Debug for Func<'a> { impl<'a> std::fmt::Debug for Func<'a> {
@ -15,9 +17,9 @@ impl<'a> std::fmt::Debug for Func<'a> {
pub fn func<'a, F>(f: F) -> Func<'a> pub fn func<'a, F>(f: F) -> Func<'a>
where where
F: 'a + Fn(&Tensor) -> Result<Tensor> + Send, F: 'a + Fn(&Tensor) -> Result<Tensor> + Send + Sync,
{ {
Func { f: Box::new(f) } Func { f: Arc::new(f) }
} }
impl<'a> super::Module for Func<'a> { impl<'a> super::Module for Func<'a> {
@ -29,8 +31,8 @@ impl<'a> super::Module for Func<'a> {
impl<'a> Func<'a> { impl<'a> Func<'a> {
pub fn new<F>(f: F) -> Self pub fn new<F>(f: F) -> Self
where where
F: 'a + Fn(&Tensor) -> Result<Tensor> + Send, F: 'a + Fn(&Tensor) -> Result<Tensor> + Send + Sync,
{ {
Self { f: Box::new(f) } Self { f: Arc::new(f) }
} }
} }

View File

@ -44,7 +44,7 @@ impl Sequential {
/// Appends a closure after all the current layers. /// Appends a closure after all the current layers.
pub fn add_fn<F>(self, f: F) -> Self pub fn add_fn<F>(self, f: F) -> Self
where where
F: 'static + Fn(&Tensor) -> Result<Tensor> + Send, F: 'static + Fn(&Tensor) -> Result<Tensor> + Send + Sync,
{ {
self.add(super::func(f)) self.add(super::func(f))
} }