refactor: simplify artifact directory handling and update optimizer configuration

- Added `reset_artifact_dir` helper function in `resnet_finetune` for cleaner artifact management.
- Replaced `AdamConfig` with `AdamWConfig` for optimizer setup.
- Removed redundant imports and allowed directives for improved code clarity.

Author: crutcher

crates/bimm/src/models/resnet/bottleneck_block.rs

@@ -522,6 +522,7 @@ impl<B: Backend> BottleneckBlock<B> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use bimm_contracts::assert_shape_contract;
     use burn::backend::NdArray;
     use burn::nn::activation::ActivationConfig;
 
@@ -571,7 +572,6 @@ mod tests {
 
     #[test]
     fn test_basic_block_forward_same_channels_no_downsample_autodiff() {
-        use bimm_contracts::assert_shape_contract;
         use burn::backend::{Autodiff, Wgpu};
         type B = Autodiff<Wgpu>;
 
@@ -603,7 +603,6 @@ mod tests {
 
     #[test]
     fn test_basic_block_forward_downsample_drop_block_drop_path_autodiff() {
-        use bimm_contracts::assert_shape_contract;
         use burn::backend::{Autodiff, Wgpu};
         type B = Autodiff<Wgpu>;
 

examples/resnet_finetune/src/main.rs

@@ -1,4 +1,4 @@
-#![allow(dead_code, unused)]
+#![allow(dead_code)]
 #![recursion_limit = "256"]
 
 extern crate core;
@@ -8,18 +8,16 @@ mod dataset;
 use crate::data::{ClassificationBatch, ClassificationBatcher};
 use crate::dataset::{CLASSES, PlanetLoader, download};
 use bimm::cache::disk::DiskCacheConfig;
-use bimm::models::resnet::{PREFAB_RESNET_MAP, ResNet, ResNetContractConfig};
+use bimm::models::resnet::{PREFAB_RESNET_MAP, ResNet};
 use burn::backend::Autodiff;
 use burn::config::Config;
 use burn::data::dataloader::DataLoaderBuilder;
 use burn::data::dataset::transform::ShuffledDataset;
 use burn::data::dataset::vision::ImageFolderDataset;
 use burn::module::Module;
-use burn::nn::PReluConfig;
 use burn::nn::activation::ActivationConfig;
 use burn::nn::loss::BinaryCrossEntropyLossConfig;
-use burn::optim::AdamConfig;
-use burn::optim::decay::WeightDecayConfig;
+use burn::optim::AdamWConfig;
 use burn::prelude::{Int, Tensor};
 use burn::record::CompactRecorder;
 use burn::tensor::backend::{AutodiffBackend, Backend};
@@ -152,6 +150,11 @@ fn main() -> anyhow::Result<()> {
     return train::<Autodiff<burn::backend::Metal>>(&args);
 }
 
+pub fn reset_artifact_dir(artifact_dir: &str) -> anyhow::Result<()> {
+    std::fs::remove_dir_all(artifact_dir)?;
+    std::fs::create_dir_all(artifact_dir).map_err(|e| anyhow::anyhow!(e))
+}
+
 #[must_use]
 pub fn train<B: AutodiffBackend>(args: &Args) -> anyhow::Result<()> {
     let device: B::Device = Default::default();
@@ -189,8 +192,7 @@ pub fn train<B: AutodiffBackend>(args: &Args) -> anyhow::Result<()> {
 
     // Remove existing artifacts before to get an accurate learner summary
     let artifact_dir: &str = args.artifact_dir.as_ref();
-    std::fs::remove_dir_all(artifact_dir);
-    std::fs::create_dir_all(artifact_dir).expect("Failed to create artifacts directory");
+    reset_artifact_dir(artifact_dir)?;
 
     B::seed(&device, args.seed);
 
@@ -224,8 +226,9 @@ pub fn train<B: AutodiffBackend>(args: &Args) -> anyhow::Result<()> {
         resnet: model,
     };
 
-    let optimizer = AdamConfig::new()
-        .with_weight_decay(Some(WeightDecayConfig::new(args.weight_decay)))
+    let optimizer = AdamWConfig::new()
+        .with_cautious_weight_decay(true)
+        .with_weight_decay(args.weight_decay)
         .init();
 
     LogConfig {