From 1ff4bd4fbcb926586778c39399d3abe5302ff783 Mon Sep 17 00:00:00 2001
From: Joshua Rosenkranz <joshua.rosenkranz@gmail.com>
Date: Thu, 28 Mar 2024 09:35:23 -0400
Subject: [PATCH] remove tp from embedding for calico with tied heads; will
 address fixing tp and tied heads in later PR

---
 fms_extras/models/calico.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/fms_extras/models/calico.py b/fms_extras/models/calico.py
index 9b7f3b1..6071457 100644
--- a/fms_extras/models/calico.py
+++ b/fms_extras/models/calico.py
@@ -185,7 +185,7 @@ def __init__(
         self.pad_id = self.config.pad_id
         self.max_expected_seq_len = self.config.max_expected_seq_len
 
-        shared = WordEmbedding(
+        self.shared = WordEmbedding(
             self.config.src_vocab_size,
             self.config.emb_dim,
             padding_idx=self.config.pad_id,
@@ -194,7 +194,6 @@ def __init__(
             tie_weights=True,
             bias=False,
         )
-        self.shared = self.distributed_strategy.distribute_module(shared)
 
         self.rot_emb = RotaryEmbedding(
             dim=self.config.emb_dim // self.config.nheads,