From 34e363e987c5c17adacc07c0cc2300f6f09dcb34 Mon Sep 17 00:00:00 2001
From: Sam Denton <106690182+sam-scale@users.noreply.github.com>
Date: Sat, 9 Mar 2024 23:48:05 -0800
Subject: [PATCH] Bug in control.py (#18)

If `model` has an attribute called `model`, then you should get the layers via `model.model.layers` not `model.layers`
---
 repeng/control.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/repeng/control.py b/repeng/control.py
index 768f18c..56d92d1 100644
--- a/repeng/control.py
+++ b/repeng/control.py
@@ -204,7 +204,7 @@ def model_layer_list(model: ControlModel | PreTrainedModel) -> torch.nn.ModuleLi
         model = model.model
 
     if hasattr(model, "model"):  # mistral-like
-        return model.layers
+        return model.model.layers
     elif hasattr(model, "transformer"):  # gpt-2-like
         return model.transformer.h
     else: