From 274bf11633f609d92729e11e16618778bed4b868 Mon Sep 17 00:00:00 2001
From: Laurent Mazare <laurent.mazare@gmail.com>
Date: Mon, 12 Feb 2024 10:26:56 +0100
Subject: [PATCH] Support defaultdict in PyTorch checkpoints. (#1696)

* Support defaultdict in PyTorch checkpoints.

* Fix clippy lint.
---
 candle-core/src/pickle.rs | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/candle-core/src/pickle.rs b/candle-core/src/pickle.rs
index f6d80830..e3f1f81d 100644
--- a/candle-core/src/pickle.rs
+++ b/candle-core/src/pickle.rs
@@ -350,8 +350,10 @@ impl Stack {
                 module_name,
                 class_name,
             } => {
-                if module_name == "collections" && class_name == "OrderedDict" {
-                    // TODO: have a separate ordered dict.
+                if module_name == "collections"
+                    && (class_name == "OrderedDict" || class_name == "defaultdict")
+                {
+                    // TODO: have a separate ordered dict and a separate default dict.
                     Some(Object::Dict(vec![]))
                 } else {
                     None