How should I write a key remap for an enum type?

[zemelLeong]

I'm getting an error with my current code.

#[derive(Module, Debug)]
pub enum Layers<B: Backend> {
    Conv2d(Conv2d<B>),
    BatchNorm(BatchNorm<B, 2>),
    ConvTranspose2d(ConvTranspose2d<B>),
}
#[derive(Module, Debug)]
pub struct DownSample<B: Backend> {
    layers: Vec<Layers<B>>,
}
LoadArgs::new(model_path.into())
  .with_key_remap("downsample.0", "downsample.layers.0.Conv2d")
  .with_key_remap("downsample.1", "downsample.layers.1.BatchNorm")

This is apart of key remap info.

---
Original Key: model.layer1.0.downsample.0.weight
Remapped Key: model.layer1.layers.0.downsample.layers.0.Conv2d.weight
Shape: [64, 64, 1, 1]
Dtype: F32
---
Original Key: model.layer1.0.downsample.1.bias
Remapped Key: model.layer1.layers.0.downsample.layers.1.BatchNorm.bias
Shape: [64]
Dtype: F32
---
Original Key: model.layer1.0.downsample.1.num_batches_tracked
Remapped Key: model.layer1.layers.0.downsample.layers.1.BatchNorm.num_batches_tracked
Shape: []
Dtype: I64
---
Original Key: model.layer1.0.downsample.1.running_mean
Remapped Key: model.layer1.layers.0.downsample.layers.1.BatchNorm.running_mean
Shape: [64]
Dtype: F32
---
Original Key: model.layer1.0.downsample.1.running_var
Remapped Key: model.layer1.layers.0.downsample.layers.1.BatchNorm.running_var
Shape: [64]
Dtype: F32
---
Original Key: model.layer1.0.downsample.1.weight
Remapped Key: model.layer1.layers.0.downsample.layers.1.BatchNorm.weight
Shape: [64]
Dtype: F32
---

thread 'main' panicked at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-import\src\pytorch\adapter.rs:91:10:
Failed to find 'weight' key in map
stack backtrace:
   0: std::panicking::begin_panic_handler
             at /rustc/82e1608dfa6e0b5569232559e3d385fea5a93112/library\std\src\panicking.rs:645
   1: core::panicking::panic_fmt
             at /rustc/82e1608dfa6e0b5569232559e3d385fea5a93112/library\core\src\panicking.rs:72
   2: core::panicking::panic_display
             at /rustc/82e1608dfa6e0b5569232559e3d385fea5a93112/library\core\src\panicking.rs:178
   3: core::panicking::panic_str
             at /rustc/82e1608dfa6e0b5569232559e3d385fea5a93112/library\core\src\panicking.rs:152
   4: core::option::expect_failed
             at /rustc/82e1608dfa6e0b5569232559e3d385fea5a93112/library\core\src\option.rs:1985
   5: enum2$<core::option::Option<enum2$<burn_core::record::serde::data::NestedValue> > >::expect<enum2$<burn_core::record::serde::data::NestedValue> >
             at /rustc/82e1608dfa6e0b5569232559e3d385fea5a93112\library\core\src\option.rs:894
   6: burn_import::pytorch::adapter::rename_weight_bias
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-import\src\pytorch\adapter.rs:89
   7: burn_import::pytorch::adapter::impl$0::adapt_batch_norm<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-import\src\pytorch\adapter.rs:59
   8: burn_core::record::serde::adapter::BurnModuleAdapter::adapt<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> > >
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\adapter.rs:10
   9: burn_core::record::serde::de::impl$1::deserialize_struct<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,burn_core::nn::norm::batch::_::impl$0::deserialize::__Visitor<b
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:63
  10: burn_core::nn::norm::batch::_::impl$0::deserialize<burn_ndarray::backend::NdArray<f64>,2,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::Deserializer<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::setting
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\nn\norm\batch.rs:26
  11: serde::de::impl$5::deserialize<burn_core::nn::norm::batch::BatchNormRecordItem<burn_ndarray::backend::NdArray<f64>,2,burn_core::record::settings::FullPrecisionSettings>,burn_core::record::serde::de::Deserializer<burn_import::pytorch::adapter::PyTorchAdapt
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:794
  12: burn_core::record::serde::de::impl$8::newtype_variant_seed<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,core::marker::PhantomData<burn_core::nn::norm::batch::BatchNo
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:529
  13: serde::de::VariantAccess::newtype_variant<burn_core::record::serde::de::ProbeEnumAccess<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> > >,burn_core::nn::norm::batch::Ba
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:2121
  14: table_structure_recognition::lore_detector::_::impl$0::deserialize::impl$2::visit_enum<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::ProbeEnumAccess<burn_import::pytorch::adapter::PyTo
             at .\src\lore_detector.rs:88
  15: burn_core::record::serde::de::impl$1::deserialize_enum<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,table_structure_recognition::lore_detector::_::impl$0::deserializ
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:365
  16: table_structure_recognition::lore_detector::_::impl$0::deserialize<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::Deserializer<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::r
             at .\src\lore_detector.rs:88
  17: serde::de::impl$5::deserialize<enum2$<table_structure_recognition::lore_detector::LayersRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings> >,burn_core::record::serde::de::Deserializer<burn_import::pytorch::a
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:794
  18: burn_core::record::serde::de::impl$3::next_element_seed<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,core::marker::PhantomData<enum2$<table_structure_recognition::lo
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:420
  19: serde::de::SeqAccess::next_element<burn_core::record::serde::de::VecSeqAccess<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> > >,enum2$<table_structure_recognition::lore
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:1728
  20: serde::de::impls::impl$19::deserialize::impl$0::visit_seq<enum2$<table_structure_recognition::lore_detector::LayersRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings> >,burn_core::record::serde::de::VecSeqAcc
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\impls.rs:1175
  21: burn_core::record::serde::de::impl$1::deserialize_seq<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,serde::de::impls::impl$19::deserialize::VecVisitor<enum2$<table_st
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:290
  22: serde::de::impls::impl$19::deserialize<enum2$<table_structure_recognition::lore_detector::LayersRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings> >,burn_core::record::serde::de::Deserializer<burn_import::py
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\impls.rs:1186
  23: serde::de::impl$5::deserialize<alloc::vec::Vec<enum2$<table_structure_recognition::lore_detector::LayersRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings> >,alloc::alloc::Global>,burn_core::record::serde::de
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:794
  24: burn_core::record::serde::de::impl$5::next_value_seed<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,core::marker::PhantomData<alloc::vec::Vec<enum2$<table_structure_r
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:477
  25: serde::de::MapAccess::next_value<burn_core::record::serde::de::HashMapAccess<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> > >,alloc::vec::Vec<enum2$<table_structure_re
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:1867
  26: table_structure_recognition::lore_detector::_::impl$0::deserialize::impl$2::visit_map<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::HashMapAccess<burn_import::pytorch::adapter::PyTorch
             at .\src\lore_detector.rs:105
  27: burn_core::record::serde::de::impl$1::deserialize_struct<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,table_structure_recognition::lore_detector::_::impl$0::deserial
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:90
  28: table_structure_recognition::lore_detector::_::impl$0::deserialize<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::Deserializer<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::r
             at .\src\lore_detector.rs:105
  29: serde::de::impls::impl$15::visit_some<table_structure_recognition::lore_detector::DownSampleRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings>,burn_core::record::serde::de::Deserializer<burn_import::pytorch:
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\impls.rs:916
  30: burn_core::record::serde::de::impl$1::deserialize_option<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,serde::de::impls::OptionVisitor<table_structure_recognition::lo
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:244
  31: serde::de::impls::impl$16::deserialize<table_structure_recognition::lore_detector::DownSampleRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings>,burn_core::record::serde::de::Deserializer<burn_import::pytorch
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\impls.rs:935
  32: serde::de::impl$5::deserialize<enum2$<core::option::Option<table_structure_recognition::lore_detector::DownSampleRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings> > >,burn_core::record::serde::de::Deseriali
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:794
  33: burn_core::record::serde::de::impl$5::next_value_seed<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,core::marker::PhantomData<enum2$<core::option::Option<table_struct
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:477
  34: serde::de::MapAccess::next_value<burn_core::record::serde::de::HashMapAccess<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> > >,enum2$<core::option::Option<table_structu
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:1867
  35: table_structure_recognition::lore_detector::_::impl$0::deserialize::impl$2::visit_map<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::HashMapAccess<burn_import::pytorch::adapter::PyTorch
             at .\src\lore_detector.rs:298
  36: burn_core::record::serde::de::impl$1::deserialize_struct<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,table_structure_recognition::lore_detector::_::impl$0::deserial
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:90
  37: table_structure_recognition::lore_detector::_::impl$0::deserialize<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::Deserializer<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::r
             at .\src\lore_detector.rs:298
  38: serde::de::impl$5::deserialize<table_structure_recognition::lore_detector::BasicBlockRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings>,burn_core::record::serde::de::Deserializer<burn_import::pytorch::adapte
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:794
  39: burn_core::record::serde::de::impl$3::next_element_seed<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,core::marker::PhantomData<table_structure_recognition::lore_dete
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:420
  40: serde::de::SeqAccess::next_element<burn_core::record::serde::de::VecSeqAccess<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> > >,table_structure_recognition::lore_detect
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:1728
  41: serde::de::impls::impl$19::deserialize::impl$0::visit_seq<table_structure_recognition::lore_detector::BasicBlockRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings>,burn_core::record::serde::de::VecSeqAccess<b
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\impls.rs:1175
  42: burn_core::record::serde::de::impl$1::deserialize_seq<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,serde::de::impls::impl$19::deserialize::VecVisitor<table_structure
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:290
  43: serde::de::impls::impl$19::deserialize<table_structure_recognition::lore_detector::BasicBlockRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings>,burn_core::record::serde::de::Deserializer<burn_import::pytorch
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\impls.rs:1186
  44: serde::de::impl$5::deserialize<alloc::vec::Vec<table_structure_recognition::lore_detector::BasicBlockRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings>,alloc::alloc::Global>,burn_core::record::serde::de::Des
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:794
  45: burn_core::record::serde::de::impl$5::next_value_seed<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,core::marker::PhantomData<alloc::vec::Vec<table_structure_recognit
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:477
  46: serde::de::MapAccess::next_value<burn_core::record::serde::de::HashMapAccess<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> > >,alloc::vec::Vec<table_structure_recogniti
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:1867
  47: table_structure_recognition::lore_detector::_::impl$0::deserialize::impl$2::visit_map<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::HashMapAccess<burn_import::pytorch::adapter::PyTorch
             at .\src\lore_detector.rs:12
  48: burn_core::record::serde::de::impl$1::deserialize_struct<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,table_structure_recognition::lore_detector::_::impl$0::deserial
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:90
  49: table_structure_recognition::lore_detector::_::impl$0::deserialize<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::Deserializer<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::r
             at .\src\lore_detector.rs:12
  50: serde::de::impl$5::deserialize<table_structure_recognition::lore_detector::SequentialRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings>,burn_core::record::serde::de::Deserializer<burn_import::pytorch::adapte
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:794
  51: burn_core::record::serde::de::impl$5::next_value_seed<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,core::marker::PhantomData<table_structure_recognition::lore_detect
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:477
  52: serde::de::MapAccess::next_value<burn_core::record::serde::de::HashMapAccess<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> > >,table_structure_recognition::lore_detecto
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:1867
  53: table_structure_recognition::lore_detector::_::impl$0::deserialize::impl$2::visit_map<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::HashMapAccess<burn_import::pytorch::adapter::PyTorch
             at .\src\lore_detector.rs:420
  54: burn_core::record::serde::de::impl$1::deserialize_struct<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,table_structure_recognition::lore_detector::_::impl$0::deserial
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:90
  55: table_structure_recognition::lore_detector::_::impl$0::deserialize<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::Deserializer<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::r
             at .\src\lore_detector.rs:420
  56: serde::de::impl$5::deserialize<table_structure_recognition::lore_detector::LoreDetectModelRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings>,burn_core::record::serde::de::Deserializer<burn_import::pytorch::a
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:794
  57: burn_core::record::serde::de::impl$5::next_value_seed<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,core::marker::PhantomData<table_structure_recognition::lore_detect
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:477
  58: serde::de::MapAccess::next_value<burn_core::record::serde::de::HashMapAccess<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> > >,table_structure_recognition::lore_detecto
             at C:\Users\11989\.cargo\registry\src\index.crates.io-6f17d22bba15001f\serde-1.0.198\src\de\mod.rs:1867
  59: table_structure_recognition::model_lore::_::impl$0::deserialize::impl$2::visit_map<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::HashMapAccess<burn_import::pytorch::adapter::PyTorchAda
             at .\src\model_lore.rs:13
  60: burn_core::record::serde::de::impl$1::deserialize_struct<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64> >,table_structure_recognition::model_lore::_::impl$0::deserialize
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-core\src\record\serde\de.rs:90
  61: table_structure_recognition::model_lore::_::impl$0::deserialize<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings,burn_core::record::serde::de::Deserializer<burn_import::pytorch::adapter::PyTorchAdapter<burn_core::reco
             at .\src\model_lore.rs:13
  62: burn_import::pytorch::reader::from_file<burn_core::record::settings::FullPrecisionSettings,table_structure_recognition::model_lore::LoreModelRecordItem<burn_ndarray::backend::NdArray<f64>,burn_core::record::settings::FullPrecisionSettings>,burn_ndarray::b
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-import\src\pytorch\reader.rs:82
  63: burn_import::pytorch::recorder::impl$0::load<burn_core::record::settings::FullPrecisionSettings,burn_ndarray::backend::NdArray<f64>,table_structure_recognition::model_lore::LoreModelRecord<burn_ndarray::backend::NdArray<f64> > >
             at C:\Users\11989\.cargo\git\checkouts\burn-178c6829f420dae1\2a721a9\crates\burn-import\src\pytorch\recorder.rs:47
  64: table_structure_recognition::model_lore::LoreModel<burn_ndarray::backend::NdArray<f64> >::new_with<burn_ndarray::backend::NdArray<f64> >
             at .\src\model_lore.rs:157
  65: table_structure_recognition::main
             at .\src\main.rs:13
  66: core::ops::function::FnOnce::call_once<void (*)(),tuple$<> >
             at /rustc/82e1608dfa6e0b5569232559e3d385fea5a93112\library\core\src\ops\function.rs:250
  67: core::hint::black_box
             at /rustc/82e1608dfa6e0b5569232559e3d385fea5a93112\library\core\src\hint.rs:286
note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace.
error: process didn't exit successfully: `J:/cargo-target\debug\table-structure-recognition.exe` (exit code: 101)

[zemelLeong]

It seems that no special handling is required for enums.

LoadArgs::new(model_path.into())
  .with_key_remap("downsample.0", "downsample.layers.0")
  .with_key_remap("downsample.1", "downsample.layers.1")

[laggui]

Will leave this here for future reference: https://burn.dev/book/import/pytorch-model.html#models-containing-enum-modules.

As @antimora noted:
Enum variants are not tagged, when importing a PyTorch model the names of both variants will obviously not be in the structure. So the best thing we can do is to match structure (attribute names and types). If you have the exact same structure for variants, it will match first one - limitation you should be aware of.