burn_train/learner/
error.rs1use burn_core::data::dataset::DatasetError;
2
3use crate::{
4 EventProcessorError,
5 checkpoint::CheckpointerError,
6 train::{MultiDeviceStepError, WorkerFailure, fmt_worker_failures},
7};
8
9#[derive(Debug)]
11pub enum TrainingError {
12 EventProcessor(EventProcessorError),
14 Dataset(DatasetError),
16 Workers(Vec<WorkerFailure>),
18 Checkpoint(CheckpointerError),
20}
21
22impl TrainingError {
23 pub fn is_device_poisoned(&self) -> bool {
26 match self {
27 Self::EventProcessor(err) => err.is_device_poisoned(),
28 _ => false,
29 }
30 }
31}
32
33impl core::fmt::Display for TrainingError {
34 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
35 match self {
36 Self::EventProcessor(err) => write!(f, "{err}"),
37 Self::Dataset(err) => write!(f, "Dataset error: {err}"),
38 Self::Workers(failures) => fmt_worker_failures(failures, f),
39 Self::Checkpoint(err) => write!(f, "Checkpoint error: {err}"),
40 }
41 }
42}
43
44impl core::error::Error for TrainingError {
45 fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
46 match self {
47 Self::EventProcessor(err) => Some(err),
48 Self::Dataset(err) => Some(err),
49 Self::Workers(_) => None,
50 Self::Checkpoint(err) => Some(err),
51 }
52 }
53}
54
55impl From<EventProcessorError> for TrainingError {
56 fn from(err: EventProcessorError) -> Self {
57 Self::EventProcessor(err)
58 }
59}
60
61impl From<DatasetError> for TrainingError {
62 fn from(err: DatasetError) -> Self {
63 Self::Dataset(err)
64 }
65}
66
67impl From<CheckpointerError> for TrainingError {
68 fn from(err: CheckpointerError) -> Self {
69 Self::Checkpoint(err)
70 }
71}
72
73impl From<MultiDeviceStepError> for TrainingError {
74 fn from(err: MultiDeviceStepError) -> Self {
75 match err {
76 MultiDeviceStepError::Dataset(err) => Self::Dataset(err),
77 MultiDeviceStepError::Workers(failures) => Self::Workers(failures),
78 }
79 }
80}