We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
BUTD training broken after PR #135
## 🐛 Bug

After PR #135, training of the BUTD method for image captioning is now broken.

This issue is that #135 removes the teacher forcing behavior during training. After #135, now BUTD to use beam search decoding (instead of teacher-forcing) during training, causing a lot of shape mismatch, gradient and other errors.

## Command

## To Reproduce

Steps to reproduce the behavior:

```
python tools/run.py \
    --tasks captioning --datasets coco --model butd \
    --config configs/captioning/coco/butd.yml
```

The training crashes due to various errors (such as batch size mismatch).

## Expected behavior

The training should work as expected.

## Environment

You can run the script with:
PyTorch version: 1.2.0
Is debug build: No
CUDA used to build PyTorch: 10.0.130

OS: Ubuntu 18.04.1 LTS
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
CMake version: version 3.9.4

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.0.130
GPU models and configuration: 
GPU 0: Quadro GP100
GPU 1: Quadro GP100

Nvidia driver version: 410.79
cuDNN version: Could not collect

Versions of relevant libraries:
[pip] block.bootstrap.pytorch==0.1.5
[pip] bootstrap.pytorch==0.0.13
[pip] msgpack-numpy==0.4.1
[pip] murel.bootstrap.pytorch==0.0.0
[pip] numpy==1.17.2
[pip] numpydoc==0.7.0
[pip] torch==1.2.0
[pip] torchfile==0.1.0
[pip] torchtext==0.3.1
[pip] torchvision==0.4.0
[conda] block.bootstrap.pytorch   0.1.5                     <pip>
[conda] bootstrap.pytorch         0.0.13                    <pip>
[conda] cuda90                    1.0                  h6433d27_0    pytorch
[conda] faiss-gpu                 1.4.0           py36_cuda9.0.176_1  [cuda90]  pytorch
[conda] magma-cuda90              2.3.0                         1    pytorch
[conda] mkl                       2018.0.3                      1  
[conda] mkl-fft                   1.0.0                     <pip>
[conda] mkl-include               2018.0.3                      1  
[conda] mkl-service               1.1.2            py36h17a0993_4  
[conda] mkl_fft                   1.0.2           np114py36_intel_0  [intel]  intel
[conda] mkl_random                1.0.1            py36h629b387_0    anaconda
[conda] mkldnn                    0.14.0                        0    mingfeima
[conda] murel.bootstrap.pytorch   0.0.0                     <pip>
[conda] nccl2                     1.0                           0    pytorch
[conda] torch                     1.2.0                     <pip>
[conda] torchfile                 0.1.0                     <pip>
[conda] torchtext                 0.3.1                     <pip>
[conda] torchvision               0.4.0                     <pip>


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pythia/models/butd.py`
Content:
```
1 # Copyright (c) Facebook, Inc. and its affiliates.
2 
3 import torch
4 from pythia.common.registry import registry
5 from pythia.modules.layers import ClassifierLayer
6 from pythia.models.pythia import Pythia
7 from pythia.utils.text_utils import BeamSearch, NucleusSampling
8 
9 
10 @registry.register_model("butd")
11 class BUTD(Pythia):
12     def __init__(self, config):
13         super().__init__(config)
14 
15     def build(self):
16         self._build_word_embedding()
17         self._init_feature_encoders("image")
18         self._init_feature_embeddings("image")
19         self._init_classifier()
20         self._init_extras()
21 
22     def _build_word_embedding(self):
23         self.text_processor = registry.get(self._datasets[0] + "_text_processor")
24         self.vocab = self.text_processor.vocab
25         self.vocab_size = self.vocab.get_size()
26         self.word_embedding = self.vocab.get_embedding(
27             torch.nn.Embedding, embedding_dim=self.config["embedding_dim"]
28         )
29         setattr(self, "text_embeddings_out_dim", self.config["embedding_dim"])
30 
31     def _init_classifier(self):
32         self.classifier = ClassifierLayer(
33             self.config["classifier"]["type"],
34             in_dim=self.config["classifier"]["params"]["feature_dim"],
35             out_dim=self.vocab_size,
36             **self.config["classifier"]["params"]
37         )
38 
39     def get_optimizer_parameters(self, config):
40         params = [
41             {"params": self.word_embedding.parameters()},
42             {"params": self.image_feature_embeddings_list.parameters()},
43             {"params": self.classifier.parameters()},
44             {
45                 "params": self.image_feature_encoders.parameters(),
46                 "lr": (config["optimizer_attributes"]["params"]["lr"] * 0.1),
47             },
48         ]
49         return params
50 
51     def prepare_data(self, sample_list, batch_size):
52         setattr(self, "teacher_forcing", hasattr(sample_list, "text"))
53         data = {}
54         if self.teacher_forcing:
55             caption_lengths, sort_ind = sample_list.caption_len.sort(
56                 dim=0, descending=True
57             )
58             data["decode_lengths"] = (caption_lengths - 1).tolist()
59             sample_list.text = sample_list.text[sort_ind]
60             sample_list.answers = sample_list.answers[sort_ind]
61             sample_list.image_feature_0 = sample_list.image_feature_0[sort_ind]
62             data["texts"] = sample_list.text
63             timesteps = max(data["decode_lengths"])
64             sample_list.add_field("targets", sample_list.text[:, 1:])
65         else:
66             data["texts"] = sample_list.answers.new_full(
67                 (batch_size, 1), self.vocab.SOS_INDEX, dtype=torch.long
68             )
69             timesteps = self.text_processor.max_length
70             sample_list.add_field("targets", sample_list.answers[:, 0, 1:])
71         return data, sample_list, timesteps
72 
73     def init_hidden_state(self, features):
74         h = features.new_zeros(
75             (features.size(0), self.config["classifier"]["params"]["hidden_dim"]),
76             dtype=torch.float,
77         )
78         c = features.new_zeros(
79             (features.size(0), self.config["classifier"]["params"]["hidden_dim"]),
80             dtype=torch.float,
81         )
82         return h, c
83 
84     def get_data_t(self, t, data, batch_size_t, prev_output):
85         if self.teacher_forcing:
86             # Modify batch_size for timestep t
87             batch_size_t = sum([l > t for l in data["decode_lengths"]])
88         elif prev_output is not None and self.config["inference"]["type"] == "greedy":
89             # Adding t-1 output words to data["text"] for greedy decoding
90             output_softmax = torch.log_softmax(prev_output, dim=1)
91             _, indices = torch.max(output_softmax, dim=1, keepdim=True)
92             data["texts"] = torch.cat(
93                 (data["texts"], indices.view(batch_size_t, 1)), dim=1
94             )
95 
96         # Slice data based on batch_size at timestep t
97         data["texts"] = data["texts"][:batch_size_t]
98         if "state" in data:
99             h1 = data["state"]["td_hidden"][0][:batch_size_t]
100             c1 = data["state"]["td_hidden"][1][:batch_size_t]
101             h2 = data["state"]["lm_hidden"][0][:batch_size_t]
102             c2 = data["state"]["lm_hidden"][1][:batch_size_t]
103         else:
104             h1, c1 = self.init_hidden_state(data["texts"])
105             h2, c2 = self.init_hidden_state(data["texts"])
106         data["state"] = {"td_hidden": (h1, c1), "lm_hidden": (h2, c2)}
107         registry.register("{}_lstm_state".format(h1.device), data["state"])
108 
109         return data, batch_size_t
110 
111     def forward(self, sample_list):
112         # Stores the output probabilites.
113         scores = sample_list.answers.new_ones(
114             (
115                 sample_list.answers.size(0),
116                 self.text_processor.max_length,
117                 self.vocab_size,
118             ),
119             dtype=torch.float,
120         )
121 
122         decoder = registry.get_decoder_class(self.config["inference"]["type"])(self.vocab, self.config)
123 
124         sample_list = decoder.init_batch(sample_list)
125         # batch_size = sample_list.get_batch_size()
126         batch_size = sample_list.image_feature_0.size(0)
127         data, sample_list, timesteps = self.prepare_data(sample_list, batch_size)
128         output = None
129         batch_size_t = batch_size
130         for t in range(timesteps):
131             data, batch_size_t = self.get_data_t(t, data, batch_size_t, output)
132             pi_t = data["texts"]
133             embedding = self.word_embedding(pi_t)
134             attention_feature, _ = self.process_feature_embedding(
135                 "image", sample_list, embedding[:, 0, :], batch_size_t=batch_size_t
136             )
137             output = self.classifier(attention_feature)
138             # Compute decoding
139             finish, data, batch_size_t = decoder.decode(t, data, output)
140             if finish:
141                 break
142 
143         model_output = {"scores": scores}
144         model_output["captions"] = decoder.get_result()
145 
146         return model_output
147
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pythia/models/butd.py b/pythia/models/butd.py
--- a/pythia/models/butd.py
+++ b/pythia/models/butd.py
@@ -119,28 +119,37 @@
             dtype=torch.float,
         )
 
-        decoder = registry.get_decoder_class(self.config["inference"]["type"])(self.vocab, self.config)
+        if self.config["inference"]["type"] in ["beam_search", "nucleus_sampling"]:
+            decoder = registry.get_decoder_class(self.config["inference"]["type"])(
+                self.vocab, self.config
+            )
+            sample_list = decoder.init_batch(sample_list)
 
-        sample_list = decoder.init_batch(sample_list)
-        # batch_size = sample_list.get_batch_size()
         batch_size = sample_list.image_feature_0.size(0)
         data, sample_list, timesteps = self.prepare_data(sample_list, batch_size)
         output = None
         batch_size_t = batch_size
         for t in range(timesteps):
             data, batch_size_t = self.get_data_t(t, data, batch_size_t, output)
-            pi_t = data["texts"]
+            if self.config["inference"]["type"] in ["beam_search", "nucleus_sampling"]:
+                pi_t = data["texts"]
+            else:
+                pi_t = data["texts"][:, t].unsqueeze(-1)
             embedding = self.word_embedding(pi_t)
             attention_feature, _ = self.process_feature_embedding(
                 "image", sample_list, embedding[:, 0, :], batch_size_t=batch_size_t
             )
             output = self.classifier(attention_feature)
             # Compute decoding
-            finish, data, batch_size_t = decoder.decode(t, data, output)
-            if finish:
-                break
+            if self.config["inference"]["type"] in ["beam_search", "nucleus_sampling"]:
+                finish, data, batch_size_t = decoder.decode(t, data, output)
+                if finish:
+                    break
+            else:
+                scores[:batch_size_t, t] = output
 
         model_output = {"scores": scores}
-        model_output["captions"] = decoder.get_result()
+        if self.config["inference"]["type"] in ["beam_search", "nucleus_sampling"]:
+            model_output["captions"] = decoder.get_result()
 
         return model_output

{"golden_diff": "diff --git a/pythia/models/butd.py b/pythia/models/butd.py\n--- a/pythia/models/butd.py\n+++ b/pythia/models/butd.py\n@@ -119,28 +119,37 @@\n             dtype=torch.float,\n         )\n \n-        decoder = registry.get_decoder_class(self.config[\"inference\"][\"type\"])(self.vocab, self.config)\n+        if self.config[\"inference\"][\"type\"] in [\"beam_search\", \"nucleus_sampling\"]:\n+            decoder = registry.get_decoder_class(self.config[\"inference\"][\"type\"])(\n+                self.vocab, self.config\n+            )\n+            sample_list = decoder.init_batch(sample_list)\n \n-        sample_list = decoder.init_batch(sample_list)\n-        # batch_size = sample_list.get_batch_size()\n         batch_size = sample_list.image_feature_0.size(0)\n         data, sample_list, timesteps = self.prepare_data(sample_list, batch_size)\n         output = None\n         batch_size_t = batch_size\n         for t in range(timesteps):\n             data, batch_size_t = self.get_data_t(t, data, batch_size_t, output)\n-            pi_t = data[\"texts\"]\n+            if self.config[\"inference\"][\"type\"] in [\"beam_search\", \"nucleus_sampling\"]:\n+                pi_t = data[\"texts\"]\n+            else:\n+                pi_t = data[\"texts\"][:, t].unsqueeze(-1)\n             embedding = self.word_embedding(pi_t)\n             attention_feature, _ = self.process_feature_embedding(\n                 \"image\", sample_list, embedding[:, 0, :], batch_size_t=batch_size_t\n             )\n             output = self.classifier(attention_feature)\n             # Compute decoding\n-            finish, data, batch_size_t = decoder.decode(t, data, output)\n-            if finish:\n-                break\n+            if self.config[\"inference\"][\"type\"] in [\"beam_search\", \"nucleus_sampling\"]:\n+                finish, data, batch_size_t = decoder.decode(t, data, output)\n+                if finish:\n+                    break\n+            else:\n+                scores[:batch_size_t, t] = output\n \n         model_output = {\"scores\": scores}\n-        model_output[\"captions\"] = decoder.get_result()\n+        if self.config[\"inference\"][\"type\"] in [\"beam_search\", \"nucleus_sampling\"]:\n+            model_output[\"captions\"] = decoder.get_result()\n \n         return model_output\n", "issue": "BUTD training broken after PR #135\n## \ud83d\udc1b Bug\r\n\r\nAfter PR #135, training of the BUTD method for image captioning is now broken.\r\n\r\nThis issue is that #135 removes the teacher forcing behavior during training. After #135, now BUTD to use beam search decoding (instead of teacher-forcing) during training, causing a lot of shape mismatch, gradient and other errors.\r\n\r\n## Command\r\n\r\n## To Reproduce\r\n\r\nSteps to reproduce the behavior:\r\n\r\n```\r\npython tools/run.py \\\r\n    --tasks captioning --datasets coco --model butd \\\r\n    --config configs/captioning/coco/butd.yml\r\n```\r\n\r\nThe training crashes due to various errors (such as batch size mismatch).\r\n\r\n## Expected behavior\r\n\r\nThe training should work as expected.\r\n\r\n## Environment\r\n\r\nYou can run the script with:\r\nPyTorch version: 1.2.0\r\nIs debug build: No\r\nCUDA used to build PyTorch: 10.0.130\r\n\r\nOS: Ubuntu 18.04.1 LTS\r\nGCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0\r\nCMake version: version 3.9.4\r\n\r\nPython version: 3.6\r\nIs CUDA available: Yes\r\nCUDA runtime version: 10.0.130\r\nGPU models and configuration: \r\nGPU 0: Quadro GP100\r\nGPU 1: Quadro GP100\r\n\r\nNvidia driver version: 410.79\r\ncuDNN version: Could not collect\r\n\r\nVersions of relevant libraries:\r\n[pip] block.bootstrap.pytorch==0.1.5\r\n[pip] bootstrap.pytorch==0.0.13\r\n[pip] msgpack-numpy==0.4.1\r\n[pip] murel.bootstrap.pytorch==0.0.0\r\n[pip] numpy==1.17.2\r\n[pip] numpydoc==0.7.0\r\n[pip] torch==1.2.0\r\n[pip] torchfile==0.1.0\r\n[pip] torchtext==0.3.1\r\n[pip] torchvision==0.4.0\r\n[conda] block.bootstrap.pytorch   0.1.5                     <pip>\r\n[conda] bootstrap.pytorch         0.0.13                    <pip>\r\n[conda] cuda90                    1.0                  h6433d27_0    pytorch\r\n[conda] faiss-gpu                 1.4.0           py36_cuda9.0.176_1  [cuda90]  pytorch\r\n[conda] magma-cuda90              2.3.0                         1    pytorch\r\n[conda] mkl                       2018.0.3                      1  \r\n[conda] mkl-fft                   1.0.0                     <pip>\r\n[conda] mkl-include               2018.0.3                      1  \r\n[conda] mkl-service               1.1.2            py36h17a0993_4  \r\n[conda] mkl_fft                   1.0.2           np114py36_intel_0  [intel]  intel\r\n[conda] mkl_random                1.0.1            py36h629b387_0    anaconda\r\n[conda] mkldnn                    0.14.0                        0    mingfeima\r\n[conda] murel.bootstrap.pytorch   0.0.0                     <pip>\r\n[conda] nccl2                     1.0                           0    pytorch\r\n[conda] torch                     1.2.0                     <pip>\r\n[conda] torchfile                 0.1.0                     <pip>\r\n[conda] torchtext                 0.3.1                     <pip>\r\n[conda] torchvision               0.4.0                     <pip>\r\n\n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n\nimport torch\nfrom pythia.common.registry import registry\nfrom pythia.modules.layers import ClassifierLayer\nfrom pythia.models.pythia import Pythia\nfrom pythia.utils.text_utils import BeamSearch, NucleusSampling\n\n\n@registry.register_model(\"butd\")\nclass BUTD(Pythia):\n    def __init__(self, config):\n        super().__init__(config)\n\n    def build(self):\n        self._build_word_embedding()\n        self._init_feature_encoders(\"image\")\n        self._init_feature_embeddings(\"image\")\n        self._init_classifier()\n        self._init_extras()\n\n    def _build_word_embedding(self):\n        self.text_processor = registry.get(self._datasets[0] + \"_text_processor\")\n        self.vocab = self.text_processor.vocab\n        self.vocab_size = self.vocab.get_size()\n        self.word_embedding = self.vocab.get_embedding(\n            torch.nn.Embedding, embedding_dim=self.config[\"embedding_dim\"]\n        )\n        setattr(self, \"text_embeddings_out_dim\", self.config[\"embedding_dim\"])\n\n    def _init_classifier(self):\n        self.classifier = ClassifierLayer(\n            self.config[\"classifier\"][\"type\"],\n            in_dim=self.config[\"classifier\"][\"params\"][\"feature_dim\"],\n            out_dim=self.vocab_size,\n            **self.config[\"classifier\"][\"params\"]\n        )\n\n    def get_optimizer_parameters(self, config):\n        params = [\n            {\"params\": self.word_embedding.parameters()},\n            {\"params\": self.image_feature_embeddings_list.parameters()},\n            {\"params\": self.classifier.parameters()},\n            {\n                \"params\": self.image_feature_encoders.parameters(),\n                \"lr\": (config[\"optimizer_attributes\"][\"params\"][\"lr\"] * 0.1),\n            },\n        ]\n        return params\n\n    def prepare_data(self, sample_list, batch_size):\n        setattr(self, \"teacher_forcing\", hasattr(sample_list, \"text\"))\n        data = {}\n        if self.teacher_forcing:\n            caption_lengths, sort_ind = sample_list.caption_len.sort(\n                dim=0, descending=True\n            )\n            data[\"decode_lengths\"] = (caption_lengths - 1).tolist()\n            sample_list.text = sample_list.text[sort_ind]\n            sample_list.answers = sample_list.answers[sort_ind]\n            sample_list.image_feature_0 = sample_list.image_feature_0[sort_ind]\n            data[\"texts\"] = sample_list.text\n            timesteps = max(data[\"decode_lengths\"])\n            sample_list.add_field(\"targets\", sample_list.text[:, 1:])\n        else:\n            data[\"texts\"] = sample_list.answers.new_full(\n                (batch_size, 1), self.vocab.SOS_INDEX, dtype=torch.long\n            )\n            timesteps = self.text_processor.max_length\n            sample_list.add_field(\"targets\", sample_list.answers[:, 0, 1:])\n        return data, sample_list, timesteps\n\n    def init_hidden_state(self, features):\n        h = features.new_zeros(\n            (features.size(0), self.config[\"classifier\"][\"params\"][\"hidden_dim\"]),\n            dtype=torch.float,\n        )\n        c = features.new_zeros(\n            (features.size(0), self.config[\"classifier\"][\"params\"][\"hidden_dim\"]),\n            dtype=torch.float,\n        )\n        return h, c\n\n    def get_data_t(self, t, data, batch_size_t, prev_output):\n        if self.teacher_forcing:\n            # Modify batch_size for timestep t\n            batch_size_t = sum([l > t for l in data[\"decode_lengths\"]])\n        elif prev_output is not None and self.config[\"inference\"][\"type\"] == \"greedy\":\n            # Adding t-1 output words to data[\"text\"] for greedy decoding\n            output_softmax = torch.log_softmax(prev_output, dim=1)\n            _, indices = torch.max(output_softmax, dim=1, keepdim=True)\n            data[\"texts\"] = torch.cat(\n                (data[\"texts\"], indices.view(batch_size_t, 1)), dim=1\n            )\n\n        # Slice data based on batch_size at timestep t\n        data[\"texts\"] = data[\"texts\"][:batch_size_t]\n        if \"state\" in data:\n            h1 = data[\"state\"][\"td_hidden\"][0][:batch_size_t]\n            c1 = data[\"state\"][\"td_hidden\"][1][:batch_size_t]\n            h2 = data[\"state\"][\"lm_hidden\"][0][:batch_size_t]\n            c2 = data[\"state\"][\"lm_hidden\"][1][:batch_size_t]\n        else:\n            h1, c1 = self.init_hidden_state(data[\"texts\"])\n            h2, c2 = self.init_hidden_state(data[\"texts\"])\n        data[\"state\"] = {\"td_hidden\": (h1, c1), \"lm_hidden\": (h2, c2)}\n        registry.register(\"{}_lstm_state\".format(h1.device), data[\"state\"])\n\n        return data, batch_size_t\n\n    def forward(self, sample_list):\n        # Stores the output probabilites.\n        scores = sample_list.answers.new_ones(\n            (\n                sample_list.answers.size(0),\n                self.text_processor.max_length,\n                self.vocab_size,\n            ),\n            dtype=torch.float,\n        )\n\n        decoder = registry.get_decoder_class(self.config[\"inference\"][\"type\"])(self.vocab, self.config)\n\n        sample_list = decoder.init_batch(sample_list)\n        # batch_size = sample_list.get_batch_size()\n        batch_size = sample_list.image_feature_0.size(0)\n        data, sample_list, timesteps = self.prepare_data(sample_list, batch_size)\n        output = None\n        batch_size_t = batch_size\n        for t in range(timesteps):\n            data, batch_size_t = self.get_data_t(t, data, batch_size_t, output)\n            pi_t = data[\"texts\"]\n            embedding = self.word_embedding(pi_t)\n            attention_feature, _ = self.process_feature_embedding(\n                \"image\", sample_list, embedding[:, 0, :], batch_size_t=batch_size_t\n            )\n            output = self.classifier(attention_feature)\n            # Compute decoding\n            finish, data, batch_size_t = decoder.decode(t, data, output)\n            if finish:\n                break\n\n        model_output = {\"scores\": scores}\n        model_output[\"captions\"] = decoder.get_result()\n\n        return model_output\n", "path": "pythia/models/butd.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n\nimport torch\nfrom pythia.common.registry import registry\nfrom pythia.modules.layers import ClassifierLayer\nfrom pythia.models.pythia import Pythia\nfrom pythia.utils.text_utils import BeamSearch, NucleusSampling\n\n\n@registry.register_model(\"butd\")\nclass BUTD(Pythia):\n    def __init__(self, config):\n        super().__init__(config)\n\n    def build(self):\n        self._build_word_embedding()\n        self._init_feature_encoders(\"image\")\n        self._init_feature_embeddings(\"image\")\n        self._init_classifier()\n        self._init_extras()\n\n    def _build_word_embedding(self):\n        self.text_processor = registry.get(self._datasets[0] + \"_text_processor\")\n        self.vocab = self.text_processor.vocab\n        self.vocab_size = self.vocab.get_size()\n        self.word_embedding = self.vocab.get_embedding(\n            torch.nn.Embedding, embedding_dim=self.config[\"embedding_dim\"]\n        )\n        setattr(self, \"text_embeddings_out_dim\", self.config[\"embedding_dim\"])\n\n    def _init_classifier(self):\n        self.classifier = ClassifierLayer(\n            self.config[\"classifier\"][\"type\"],\n            in_dim=self.config[\"classifier\"][\"params\"][\"feature_dim\"],\n            out_dim=self.vocab_size,\n            **self.config[\"classifier\"][\"params\"]\n        )\n\n    def get_optimizer_parameters(self, config):\n        params = [\n            {\"params\": self.word_embedding.parameters()},\n            {\"params\": self.image_feature_embeddings_list.parameters()},\n            {\"params\": self.classifier.parameters()},\n            {\n                \"params\": self.image_feature_encoders.parameters(),\n                \"lr\": (config[\"optimizer_attributes\"][\"params\"][\"lr\"] * 0.1),\n            },\n        ]\n        return params\n\n    def prepare_data(self, sample_list, batch_size):\n        setattr(self, \"teacher_forcing\", hasattr(sample_list, \"text\"))\n        data = {}\n        if self.teacher_forcing:\n            caption_lengths, sort_ind = sample_list.caption_len.sort(\n                dim=0, descending=True\n            )\n            data[\"decode_lengths\"] = (caption_lengths - 1).tolist()\n            sample_list.text = sample_list.text[sort_ind]\n            sample_list.answers = sample_list.answers[sort_ind]\n            sample_list.image_feature_0 = sample_list.image_feature_0[sort_ind]\n            data[\"texts\"] = sample_list.text\n            timesteps = max(data[\"decode_lengths\"])\n            sample_list.add_field(\"targets\", sample_list.text[:, 1:])\n        else:\n            data[\"texts\"] = sample_list.answers.new_full(\n                (batch_size, 1), self.vocab.SOS_INDEX, dtype=torch.long\n            )\n            timesteps = self.text_processor.max_length\n            sample_list.add_field(\"targets\", sample_list.answers[:, 0, 1:])\n        return data, sample_list, timesteps\n\n    def init_hidden_state(self, features):\n        h = features.new_zeros(\n            (features.size(0), self.config[\"classifier\"][\"params\"][\"hidden_dim\"]),\n            dtype=torch.float,\n        )\n        c = features.new_zeros(\n            (features.size(0), self.config[\"classifier\"][\"params\"][\"hidden_dim\"]),\n            dtype=torch.float,\n        )\n        return h, c\n\n    def get_data_t(self, t, data, batch_size_t, prev_output):\n        if self.teacher_forcing:\n            # Modify batch_size for timestep t\n            batch_size_t = sum([l > t for l in data[\"decode_lengths\"]])\n        elif prev_output is not None and self.config[\"inference\"][\"type\"] == \"greedy\":\n            # Adding t-1 output words to data[\"text\"] for greedy decoding\n            output_softmax = torch.log_softmax(prev_output, dim=1)\n            _, indices = torch.max(output_softmax, dim=1, keepdim=True)\n            data[\"texts\"] = torch.cat(\n                (data[\"texts\"], indices.view(batch_size_t, 1)), dim=1\n            )\n\n        # Slice data based on batch_size at timestep t\n        data[\"texts\"] = data[\"texts\"][:batch_size_t]\n        if \"state\" in data:\n            h1 = data[\"state\"][\"td_hidden\"][0][:batch_size_t]\n            c1 = data[\"state\"][\"td_hidden\"][1][:batch_size_t]\n            h2 = data[\"state\"][\"lm_hidden\"][0][:batch_size_t]\n            c2 = data[\"state\"][\"lm_hidden\"][1][:batch_size_t]\n        else:\n            h1, c1 = self.init_hidden_state(data[\"texts\"])\n            h2, c2 = self.init_hidden_state(data[\"texts\"])\n        data[\"state\"] = {\"td_hidden\": (h1, c1), \"lm_hidden\": (h2, c2)}\n        registry.register(\"{}_lstm_state\".format(h1.device), data[\"state\"])\n\n        return data, batch_size_t\n\n    def forward(self, sample_list):\n        # Stores the output probabilites.\n        scores = sample_list.answers.new_ones(\n            (\n                sample_list.answers.size(0),\n                self.text_processor.max_length,\n                self.vocab_size,\n            ),\n            dtype=torch.float,\n        )\n\n        if self.config[\"inference\"][\"type\"] in [\"beam_search\", \"nucleus_sampling\"]:\n            decoder = registry.get_decoder_class(self.config[\"inference\"][\"type\"])(\n                self.vocab, self.config\n            )\n            sample_list = decoder.init_batch(sample_list)\n\n        batch_size = sample_list.image_feature_0.size(0)\n        data, sample_list, timesteps = self.prepare_data(sample_list, batch_size)\n        output = None\n        batch_size_t = batch_size\n        for t in range(timesteps):\n            data, batch_size_t = self.get_data_t(t, data, batch_size_t, output)\n            if self.config[\"inference\"][\"type\"] in [\"beam_search\", \"nucleus_sampling\"]:\n                pi_t = data[\"texts\"]\n            else:\n                pi_t = data[\"texts\"][:, t].unsqueeze(-1)\n            embedding = self.word_embedding(pi_t)\n            attention_feature, _ = self.process_feature_embedding(\n                \"image\", sample_list, embedding[:, 0, :], batch_size_t=batch_size_t\n            )\n            output = self.classifier(attention_feature)\n            # Compute decoding\n            if self.config[\"inference\"][\"type\"] in [\"beam_search\", \"nucleus_sampling\"]:\n                finish, data, batch_size_t = decoder.decode(t, data, output)\n                if finish:\n                    break\n            else:\n                scores[:batch_size_t, t] = output\n\n        model_output = {\"scores\": scores}\n        if self.config[\"inference\"][\"type\"] in [\"beam_search\", \"nucleus_sampling\"]:\n            model_output[\"captions\"] = decoder.get_result()\n\n        return model_output\n", "path": "pythia/models/butd.py"}]}