From 5fce525f9058c82784cc76b92a906a2d9b882361 Mon Sep 17 00:00:00 2001
From: Marc Sun <57196510+SunMarc@users.noreply.github.com>
Date: Fri, 1 Mar 2024 10:32:31 -0500
Subject: [PATCH] Fix edge case in infer_auto_device_map when dealing with
 buffers (#2511)

* fix buffer

* style
---
 src/accelerate/utils/modeling.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 7ad0adee801..8ca67fa15a1 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -1153,7 +1153,11 @@ def infer_auto_device_map(
         # Case 1 -> We're too big!
         if current_max_size is not None and current_memory_used + module_size > current_max_size:
             # Split or not split?
-            modules_children = [] if isinstance(module, nn.Parameter) else list(module.named_children())
+            modules_children = (
+                []
+                if isinstance(module, nn.Parameter) or isinstance(module, torch.Tensor)
+                else list(module.named_children())
+            )
             if verbose:
                 print(
                     f"Not enough space on {devices[current_device]} to put {name} (space available "