speaker.py

...
batch_size = inputs.shape[0]
input_embs = sender_embedding(inputs)
inputs = input_embs.view(batch_size, num_digits * embedding_size_sender)
hx = torch.zeros(batch_size, num_lstm_sender)
cx = torch.zeros(batch_size, num_lstm_sender)
for num in range(num_binary_messages):
  hx, cx = sender_cell(inputs, (hx, cx))
  output = sender_project(hx)
  pre_logits = sender_out(output)
  sample = gumbel_softmax(pre_logits, temperature[num])
  ...