Python chainer.functions.concat() Examples

def __call__(self, s, xs):
        """Calculate all hidden states and cell states.
        Args:
            s  (~chainer.Variable or None): Initial (hidden & cell) states. If ``None``
                is specified zero-vector is used.
            xs (list of ~chianer.Variable): List of input sequences.
                Each element ``xs[i]`` is a :class:`chainer.Variable` holding
                a sequence.
        Return:
            (hy,cy): a pair of hidden and cell states at the end of the sequence,
            ys: a hidden state sequence at the last layer
        """
        if len(xs) > 1:
            sections = np.cumsum(np.array([len(x) for x in xs[:-1]], dtype=np.int32))
            xs = F.split_axis(self.embed(F.concat(xs, axis=0)), sections, axis=0)
        else:
            xs = [ self.embed(xs[0]) ]
        if s is not None:
            hy, cy, ys = self.lstm(s[0], s[1], xs)
        else:
            hy, cy, ys = self.lstm(None, None, xs)

        return (hy,cy), ys 

def __call__(self, prev_hg, prev_he, prev_ce, x, v, r, u):
        xu = cf.concat((x, u), axis=1)
        xu = self.downsample_xu(xu)
        v = self.broadcast_v(v)
        if r.shape[2] == 1:
            r = self.broadcast_r(r)

        lstm_input = cf.concat((prev_he, prev_hg, xu, v, r), axis=1)
        gate_inputs = self.lstm(lstm_input)

        if self.use_cuda_kernel:
            next_h, next_c = CoreFunction()(gate_inputs, prev_ce)
        else:
            forget_gate_input, input_gate_input, tanh_input, output_gate_input = cf.split_axis(
                gate_inputs, 4, axis=1)

            forget_gate = cf.sigmoid(forget_gate_input)
            input_gate = cf.sigmoid(input_gate_input)
            next_c = forget_gate * prev_ce + input_gate * cf.tanh(tanh_input)
            output_gate = cf.sigmoid(output_gate_input)
            next_h = output_gate * cf.tanh(next_c)

        return next_h, next_c 

def __call__(self, prev_hg, prev_cg, prev_z, v, r, prev_u):
        v = self.broadcast_v(v)
        if r.shape[2] == 1:
            r = self.broadcast_r(r)

        lstm_input = cf.concat((prev_hg, v, r, prev_z), axis=1)
        gate_inputs = self.lstm(lstm_input)

        forget_gate_input, input_gate_input, tanh_input, output_gate_input = cf.split_axis(
            gate_inputs, 4, axis=1)

        forget_gate = cf.sigmoid(forget_gate_input)
        input_gate = cf.sigmoid(input_gate_input)
        next_c = forget_gate * prev_cg + input_gate * cf.tanh(tanh_input)
        output_gate = cf.sigmoid(output_gate_input)
        next_h = output_gate * cf.tanh(next_c)

        next_u = self.upsample_h(next_h) + prev_u

        return next_h, next_c, next_u 

def __call__(self, x):
        if self.dr:
            with chainer.using_config('train', True):
                x = F.dropout(x, self.dr)
        if self.gap:
            x = F.sum(x, axis=(2,3))
        N = x.shape[0]
        #Below code copyed from https://github.com/pfnet-research/chainer-gan-lib/blob/master/minibatch_discrimination/net.py
        feature = F.reshape(F.leaky_relu(x), (N, -1))
        m = F.reshape(self.md(feature), (N, self.B * self.C, 1))
        m0 = F.broadcast_to(m, (N, self.B * self.C, N))
        m1 = F.transpose(m0, (2, 1, 0))
        d = F.absolute(F.reshape(m0 - m1, (N, self.B, self.C, N)))
        d = F.sum(F.exp(-F.sum(d, axis=2)), axis=2) - 1
        h = F.concat([feature, d])

        h = self.l(h)
        return h 

def n_step_forward(self, x, recurrent_state):
        """Multi-step batch forward computation.

        This method sequentially applies layers as chainer.Sequential does.

        Args:
            x (list): Input sequences. Each sequence should be a variable whose
                first axis corresponds to time or a tuple of such variables.
            recurrent_state (object): Batched recurrent state. If set to None,
                it is initialized.
            output_mode (str): If set to 'concat', the output value is
                concatenated into a single large batch, which can be suitable
                for loss computation. If set to 'split', the output value is
                a list of output sequences.

        Returns:
            object: Output sequences. See the description of the `output_mode`
                argument.
            object: New batched recurrent state.
        """
        raise NotImplementedError 

def __call__(self, x, recurrent_state):
        """One-step batch forward computation.

        Args:
            x (chainer.Variable, ndarray, or tuple): One-step batched input.
            recurrent_state (object): Batched recurrent state.

        Returns:
            chainer.Variable, ndarray, or tuple: One-step batched output.
            object: New batched recurrent state.
        """
        assert isinstance(x, (chainer.Variable, self.xp.ndarray))
        return self.n_step_forward(
            split_one_step_batch_input(x),
            recurrent_state,
            output_mode='concat',
        ) 

def channelwise_inhibited(self, h):
        self.c = random.randint(0, 2)
        xp = cuda.get_array_module(h.data)
        num = h.data.shape[0]

        h = F.split_axis(h, 3, 1)
        c = F.reshape(h[self.c], (num, 16, 16))
        z = Variable(xp.zeros_like(c.data), 'AUTO')
        c = F.batch_matmul(c, z)
        c = F.reshape(c, (num, 1, 16, 16))
        hs = []
        for i, s in enumerate(h):
            if i == self.c:
                hs.append(c)
            else:
                hs.append(s)
        return F.concat(hs, 1) 

def channelwise_inhibited(self, h):
        xp = cuda.get_array_module(h.data)
        num = h.data.shape[0]

        h = F.split_axis(h, 3, 1)
        c = F.reshape(h[self.c], (num, 16, 16))
        z = Variable(xp.zeros_like(c.data), 'AUTO')
        c = F.batch_matmul(c, z)
        c = F.reshape(c, (num, 1, 16, 16))
        hs = []
        for i, s in enumerate(h):
            if i == self.c:
                hs.append(c)
            else:
                hs.append(s)
        return F.concat(hs, 1) 

def apply_to_seq(self, seq_list):
        mb_size = len(seq_list)
        mb_initial_cell, mb_initial_state = self.get_initial_states(mb_size)
        return self.nstep_lstm(mb_initial_cell, mb_initial_state, seq_list)


# class DoubleGRU(Chain):
#     def __init__(self, H, I):
#         log.info("using double GRU")
#         self.H1 = H/2
#         self.H2 = H - self.H1
#         super(DoubleGRU, self).__init__(
#             gru1 = faster_gru.GRU(self.H1, I),
#             gru2 = faster_gru.GRU(self.H2, self.H1)
#         )
#
#     def __call__(self, prev_state, inpt):
#         prev_state1, prev_state2 = F.split_axis(prev_state, (self.H1,), axis = 1)
#
#         prev_state1 = self.gru1(prev_state1, inpt)
#         prev_state2 = self.gru2(prev_state2, prev_state1)
#
#         return F.concat((prev_state1, prev_state2), axis = 1) 

def advance_state(self, previous_states, prev_y):
        current_mb_size = prev_y.data.shape[0]
        assert self.mb_size is None or current_mb_size <= self.mb_size

        if current_mb_size < len(previous_states[0].data):
            truncated_states = [None] * len(previous_states)
            for num_state in six.moves.range(len(previous_states)):
                truncated_states[num_state], _ = F.split_axis(
                    previous_states[num_state], (current_mb_size,), 0)
            previous_states = tuple(truncated_states)

        output_state = previous_states[-1]
        if self.decoder_chain.use_goto_attention:
            ci, attn = self.compute_ctxt(output_state, prev_y)
        else:
            ci, attn = self.compute_ctxt(output_state)
        concatenated = F.concat((prev_y, ci))

        new_states = self.decoder_chain.gru(previous_states, concatenated)
        return new_states, concatenated, attn 

def compute_logits(self, seq_list, encoded_input, mask_input):
        mb_size = len(seq_list)
        max_length_1 = max(len(x) for x in seq_list)
        x, mask = self.make_batch(seq_list)
        
#         print "padded_data", x
#         print "mask", mask
        
        assert self.xp.all(mask_input == self.xp.broadcast_to(mask_input[:,0:1,0:1,:], mask_input.shape))
        
        encoded = self.emb(x)
        encoded += self.get_pos_vect(mb_size, max_length_1)
        
        if self.dropout is not None:
            encoded = F.dropout(encoded, self.dropout)
        
        bos_plus_encoded = F.concat((F.broadcast_to(self.bos_encoding, (mb_size, 1, self.d_model)), encoded), axis=1)
        
        cross_mask = self.xp.broadcast_to(mask_input[:,0:1,0:1,:], (mask_input.shape[0], self.n_heads, bos_plus_encoded.data.shape[1], mask_input.shape[3]))
        
        final_layer =  self.encoding_layers(bos_plus_encoded, encoded_input, mask, cross_mask)
        logits = apply_linear_layer_to_last_dims(final_layer, self.logits_layer)
        return logits 

def __call__(self, x):
        br0 = self.Branch_0.Conv2d_1x1(x)

        br1 = self.Branch_1.Conv2d_0a_1x1(x)
        br1 = self.Branch_1.Conv2d_0b_3x3(br1)

        br2 = self.Branch_2.Conv2d_0a_1x1(x)
        br2 = self.Branch_2.Conv2d_0b_3x3(br2)
        br2 = self.Branch_2.Conv2d_0c_3x3(br2)

        mixed = F.concat((br0, br1, br2), axis=1)

        lazy_init_conv_to_join(self, x)
        up = self.Conv2d_1x1(mixed)

        x += self.scale * up
        if self.activation_fn is not None:
            x = self.activation_fn(x)
        return x 

def __call__(self, x):
        br0 = self.Branch_0.Conv2d_1x1(x)

        br1 = self.Branch_1.Conv2d_0a_1x1(x)
        br1 = self.Branch_1.Conv2d_0b_1x7(br1)
        br1 = self.Branch_1.Conv2d_0c_7x1(br1)

        mixed = F.concat((br0, br1), axis=1)

        lazy_init_conv_to_join(self, x)
        up = self.Conv2d_1x1(mixed)

        x += self.scale * up
        if self.activation_fn is not None:
            x = self.activation_fn(x)
        return x 

def __call__(self, x):
        br0 = self.Branch_0.Conv2d_1x1(x)

        br1 = self.Branch_1.Conv2d_0a_1x1(x)
        br1 = self.Branch_1.Conv2d_0b_1x3(br1)
        br1 = self.Branch_1.Conv2d_0c_3x1(br1)

        mixed = F.concat((br0, br1), axis=1)

        lazy_init_conv_to_join(self, x)
        up = self.Conv2d_1x1(mixed)

        x += self.scale * up
        if self.activation_fn is not None:
            x = self.activation_fn(x)
        return x 

def __call__(self, s, xs):
        """Calculate all hidden states and cell states.
        Args:
            s  (~chainer.Variable or None): Initial (hidden & cell) states. If ``None``
                is specified zero-vector is used.
            xs (list of ~chianer.Variable): List of input sequences.
                Each element ``xs[i]`` is a :class:`chainer.Variable` holding
                a sequence.
        Return:
            (hy,cy): a pair of hidden and cell states at the end of the sequence,
            ys: a hidden state sequence at the last layer
        """
        if len(xs) > 1:
            sections = np.cumsum(np.array([len(x) for x in xs[:-1]], dtype=np.int32))
            xs = F.split_axis(self.embed(F.concat(xs, axis=0)), sections, axis=0)
        else:
            xs = [ self.embed(xs[0]) ]
        if s is not None:
            hy, cy, ys = self.lstm(s[0], s[1], xs)
        else:
            hy, cy, ys = self.lstm(None, None, xs)

        return (hy,cy), ys 

def __call__(self, s, xs):
        """Calculate all hidden states and cell states.
        Args:
            s  (~chainer.Variable or None): Initial (hidden & cell) states. If ``None``
                is specified zero-vector is used.
            xs (list of ~chianer.Variable): List of input sequences.
                Each element ``xs[i]`` is a :class:`chainer.Variable` holding
                a sequence.
        Return:
            (hy,cy): a pair of hidden and cell states at the end of the sequence,
            ys: a hidden state sequence at the last layer
        """
        if len(xs) > 1:
            sections = np.cumsum(np.array([len(x) for x in xs[:-1]], dtype=np.int32))
            xs = F.split_axis(self.embed(F.concat(xs, axis=0)), sections, axis=0)
        else:
            xs = [ self.embed(xs[0]) ]
        if s is not None:
            hy, cy, ys = self.lstm(s[0], s[1], xs)
        else:
            hy, cy, ys = self.lstm(None, None, xs)

        return (hy,cy), ys 

def naive_call(self, fb_concat, targets, mask):
        compute_ctxt = self.attn_module.naive_call(fb_concat, mask)
        loss = None
        current_mb_size = targets[0].data.shape[0]
        assert current_mb_size == 1
        previous_states = self.gru.get_initial_states(current_mb_size)
#         previous_word = Variable(np.array([self.bos_idx] * mb_size, dtype = np.int32))
        # xp = cuda.get_array_module(self.gru.initial_state.data)
        with cuda.get_device_from_array(self.gru.initial_state.data):
            prev_y = F.broadcast_to(self.bos_embeding, (1, self.Eo))
#             previous_word = Variable(xp.array([self.bos_idx] * current_mb_size, dtype = np.int32))
        previous_word = None
        attn_list = []
        total_nb_predictions = 0
        for i in range(len(targets)):
            if previous_word is not None:  # else we are using the initial prev_y
                prev_y = self.emb(previous_word)
            ci, attn = compute_ctxt(previous_states[-1])
            concatenated = F.concat((prev_y, ci))
    #             print(concatenated.data.shape)
            new_states = self.gru(previous_states, concatenated)

            all_concatenated = F.concat((concatenated, new_states[-1]))
            logits = self.lin_o(self.maxo(all_concatenated))

            local_loss = F.softmax_cross_entropy(logits, targets[i])

            loss = local_loss if loss is None else loss + local_loss
            total_nb_predictions += 1
            previous_word = targets[i]
            previous_states = new_states
            attn_list.append(attn)

        loss = loss / total_nb_predictions
        return loss, attn_list 

def forward(self, xs, ys):
        xs = [x[::-1] for x in xs]

        eos = self.xp.array([EOS], numpy.int32)
        ys_in = [F.concat([eos, y], axis=0) for y in ys]
        ys_out = [F.concat([y, eos], axis=0) for y in ys]

        # Both xs and ys_in are lists of arrays.
        exs = sequence_embed(self.embed_x, xs)
        eys = sequence_embed(self.embed_y, ys_in)

        batch = len(xs)
        # None represents a zero vector in an encoder.
        hx, cx, _ = self.encoder(None, None, exs)
        _, _, os = self.decoder(hx, cx, eys)

        # It is faster to concatenate data before calculating loss
        # because only one matrix multiplication is called.
        concat_os = F.concat(os, axis=0)
        concat_ys_out = F.concat(ys_out, axis=0)
        loss = F.sum(F.softmax_cross_entropy(
            self.W(concat_os), concat_ys_out, reduce='no')) / batch

        chainer.report({'loss': loss.data}, self)
        n_words = concat_ys_out.shape[0]
        perp = self.xp.exp(loss.data * batch / n_words)
        chainer.report({'perp': perp}, self)
        return loss


# from https://github.com/chainer/chainer/blob/master/examples/seq2seq/seq2seq.py 

def faster_call2(self, h, x):
        h_x = F.concat((h, x), axis=1)
        z_r = F.sigmoid(self.W_r_z(h_x))

        z, r = F.split_axis(z_r, (self.n_units,), axis=1)

        hr_x = F.concat((h * r, x), axis=1)

        h_new = compute_GRU_out_2(z, h, self.W_h(hr_x))
        return h_new 

def faster_call(self, h, x):
        h_x = F.concat((h, x), axis=1)
        z_r = F.sigmoid(self.W_r_z(h_x))

        z, r = F.split_axis(z_r, (self.n_units,), axis=1)

        hr_x = F.concat((h * r, x), axis=1)
        h_bar = F.tanh(self.W_h(hr_x))

        h_new = (1 - z) * h + z * h_bar
        return h_new 

def compute_loss(self, hx, dataset_as_array_list, 
                     need_to_append_eos = True,
                    use_gumbel = False,
                    temperature = None,
                     train = True,
                    verbose = False):
        if need_to_append_eos:
            dataset_as_var_with_eos = [Variable(self.append_eos_id(a)) for a in dataset_as_array_list]
        else:
            dataset_as_var_with_eos = [Variable(a) for a in dataset_as_array_list]
            
        if need_to_append_eos:
            dataset_as_var_without_eos = [Variable(a) for a in dataset_as_array_list]
        else:
            dataset_as_var_without_eos = [Variable(a[:-1]) for a in dataset_as_array_list]
            
        nb_ex = len(dataset_as_array_list)
        dataset_as_emb_dec = [self.c_emb_dec(v) for v in dataset_as_var_without_eos]
        hx_dec, cx_dec, xs_dec = self.nstep_dec(hx, None, dataset_as_emb_dec, train = train)
        hx_initial = F.split_axis(hx.reshape(nb_ex, -1), nb_ex, axis = 0, force_tuple = True)
        logits_list = [self.lin_out(F.concat((hxi, h), axis = 0)) for hxi, h in zip(hx_initial, xs_dec)]
    
        if verbose:
            print "logits:"
            for logits in logits_list:
                print logits.data
            print
            
        if use_gumbel:
            logits_list = [(logits + self.xp.random.gumbel(size = logits.data.shape)) for logits in logits_list]
        
        if temperature is not None:
            logits_list = [logits/temperature for logits in logits_list]
        
        losses = [F.softmax_cross_entropy(logits, tgt) for logits,tgt in zip(logits_list, dataset_as_var_with_eos)]
        loss = sum(losses)/nb_ex
        return loss 

def compute_logits(self, new_states, concatenated, attn):
        new_output_state = new_states[-1]

        all_concatenated = F.concat((concatenated, new_output_state))
        logits = self.decoder_chain.lin_o(self.decoder_chain.maxo(all_concatenated))

        if self.lexicon_probability_matrix is not None:
            current_mb_size = new_output_state.data.shape[0]
            assert self.mb_size is None or current_mb_size <= self.mb_size
            lexicon_probability_matrix = self.lexicon_probability_matrix[:current_mb_size]

            # Just making sure data shape is as expected
            attn_mb_size, max_source_length_attn = attn.data.shape
            assert attn_mb_size == current_mb_size
            lex_mb_size, max_source_length_lexicon, v_size_lexicon = lexicon_probability_matrix.shape
            assert max_source_length_lexicon == max_source_length_attn
            assert logits.data.shape == (current_mb_size, v_size_lexicon)

            if self.demux:
                assert lex_mb_size == 1
                weighted_lex_probs = F.reshape(
                    matmul_constant(attn, lexicon_probability_matrix.reshape(lexicon_probability_matrix.shape[1],
                                                                             lexicon_probability_matrix.shape[2])),
                    logits.data.shape)
            else:
                assert lex_mb_size == current_mb_size

    #                 weighted_lex_probs = F.reshape(
    #                         F.batch_matmul(attn, ConstantFunction(lexicon_probability_matrix)(), transa = True),
    #                                                logits.data.shape)

                weighted_lex_probs = F.reshape(
                    batch_matmul_constant(attn, lexicon_probability_matrix, transa=True),
                    logits.data.shape)

            logits += F.log(weighted_lex_probs + self.lex_epsilon)
        return logits 

def naive_call(self, sequence, mask):

        mb_size = sequence[0].data.shape[0]
        assert mb_size == 1

        mb_initial_states_f = self.gru_f.get_initial_states(mb_size)
        mb_initial_states_b = self.gru_b.get_initial_states(mb_size)

        embedded_seq = []
        for elem in sequence:
            embedded_seq.append(self.emb(elem))

#         self.gru_f.reset_state()
        prev_states = mb_initial_states_f
        forward_seq = []
        for i, x in enumerate(embedded_seq):
            prev_states = self.gru_f(prev_states, x)
            forward_seq.append(prev_states)

#         self.gru_b.reset_state()
        prev_states = mb_initial_states_b
        backward_seq = []
        for pos, x in reversed(list(enumerate(embedded_seq))):
            prev_states = self.gru_b(prev_states, x)
            backward_seq.append(prev_states)

        assert len(backward_seq) == len(forward_seq)
        res = []
        for xf, xb in zip(forward_seq, reversed(backward_seq)):
            res.append(F.concat((xf[-1], xb[-1]), 1))

        return res 

def __call__(self, h, *h_skip):
        h = F.relu(self.b0(self.d0(h)))
        h = F.concat((h, *h_skip))
        h = F.relu(self.b1(self.c1(h)))
        h = F.relu(self.b2(self.c2(h)))

        return h 

def __call__(self, x, ss):
        # Inception (we don't need forget, as we don't backprop)
        p0 = ss[0]
        c = (chainer.configuration.config.max_perturbation - 4) * 3
        x_plus_p0 = x + p0[:, c:c + 3, :, :]
        ds = self.ins(x_plus_p0)
        ss = [F.concat((s, d), axis=1) for s, d in zip(ss, ds)]

        # Decode (forget)
        assert isinstance(x, chainer.Variable)
        ss = F.forget(lambda x_, p0_, *ss_: self.dec(x_, p0_, ss_), x, p0, *ss)

        return ss 

def __call__(self, x, v):
        v = self.broadcast_v(v)
        resnet_in = cf.relu(self.conv1_1(x))
        residual = cf.relu(self.conv1_res(resnet_in))
        out = cf.relu(self.conv1_2(resnet_in))
        out = cf.relu(self.conv1_3(out)) + residual
        resnet_in = cf.concat((out, v), axis=1)
        residual = cf.relu(self.conv2_res(resnet_in))
        out = cf.relu(self.conv2_1(resnet_in))
        out = cf.relu(self.conv2_2(out)) + residual
        out = self.conv2_3(out)
        return out 

def __call__(self, x):
        N = x.shape[0]
        #Below code copyed from https://github.com/pfnet-research/chainer-gan-lib/blob/master/minibatch_discrimination/net.py
        feature = F.reshape(x, (N, -1))
        m = F.reshape(self.md(feature), (N, self.B * self.C, 1))
        m0 = F.broadcast_to(m, (N, self.B * self.C, N))
        m1 = F.transpose(m0, (2, 1, 0))
        d = F.absolute(F.reshape(m0 - m1, (N, self.B, self.C, N)))
        d = F.sum(F.exp(-F.sum(d, axis=2)), axis=2) - 1
        h = F.concat([feature, d])

        h = self.l(h)
        return h 

def _make_dis_input(self, input_img, label_map):
        b = F.broadcast_to(input_img[:,0,:,:], shape=label_map.shape)
        g = F.broadcast_to(input_img[:,1,:,:], shape=label_map.shape)
        r = F.broadcast_to(input_img[:,2,:,:], shape=label_map.shape)
        product_b = label_map * b
        product_g = label_map * g
        product_r = label_map * r
        dis_input = F.concat([product_b, product_g, product_r], axis=1)
        return dis_input 

def sequence_embed(embed, xs):
    x_len = [len(x) for x in xs]
    x_section = numpy.cumsum(x_len[:-1])
    ex = embed(F.concat(xs, axis=0))
    exs = F.split_axis(ex, x_section, 0)
    return exs 

def compute_loss(self, hx, dataset_as_array_list, 
                     need_to_append_eos = True,
                    use_gumbel = False,
                    temperature = None,
                     train = True,
                    verbose = False):
        if need_to_append_eos:
            dataset_as_var_with_eos = [Variable(self.append_eos_id(a)) for a in dataset_as_array_list]
        else:
            dataset_as_var_with_eos = [Variable(a) for a in dataset_as_array_list]
            
        if need_to_append_eos:
            dataset_as_var_without_eos = [Variable(a) for a in dataset_as_array_list]
        else:
            dataset_as_var_without_eos = [Variable(a[:-1]) for a in dataset_as_array_list]
            
        nb_ex = len(dataset_as_array_list)
        dataset_as_emb_dec = [self.c_emb_dec(v) for v in dataset_as_var_without_eos]
        hx_dec, cx_dec, xs_dec = self.nstep_dec(hx, None, dataset_as_emb_dec, train = train)
        hx_initial = F.split_axis(hx.reshape(nb_ex, -1), nb_ex, axis = 0, force_tuple = True)
        logits_list = [self.lin_out(F.concat((hxi, h), axis = 0)) for hxi, h in zip(hx_initial, xs_dec)]
    
        if verbose:
            print "logits:"
            for logits in logits_list:
                print logits.data
            print
            
        if use_gumbel:
            logits_list = [(logits + self.xp.random.gumbel(size = logits.data.shape)) for logits in logits_list]
        
        if temperature is not None:
            logits_list = [logits/temperature for logits in logits_list]
        
        losses = [F.softmax_cross_entropy(logits, tgt) for logits,tgt in zip(logits_list, dataset_as_var_with_eos)]
        loss = sum(losses)/nb_ex
        return loss