def forward(self, input_, hx):
        """
        Args:
            input_: A (batch, input_size) tensor containing input
                features.
            hx: A tuple (h_0, c_0), which contains the initial hidden
                and cell state, where the size of both states is
                (batch, hidden_size).
            time: The current timestep value, which is used to
                get appropriate running statistics.

        Returns:
            h_1, c_1: Tensors containing the next hidden and cell state.
        """

        h_0, c_0 = hx
        batch_size = h_0.size(0)
        bias_batch = (self.bias.unsqueeze(0)
                      .expand(batch_size, *self.bias.size()))
        wh = torch.mm(h_0, self.weight_hh)
        wh = torch.mm(h_0, self.weight_hh)
        wi = torch.mm(input_, self.weight_ih)
        bn_wh = self.bn_hh(wh)
        bn_wi = self.bn_ih(wi)
        f, i, o, g = torch.split(bn_wh + bn_wi + bias_batch,
                                 split_size=self.hidden_size, dim=1)
        c_1 = torch.sigmoid(f)*c_0 + torch.sigmoid(i)*torch.tanh(g)
        h_1 = torch.sigmoid(o) * torch.tanh(self.bn_c(c_1))
        return h_1, c_1

    def forward(self, embbedings, label):
        if self.device_id == None:
            kernel_norm = l2_norm(self.kernel, axis = 0)
            cos_theta = torch.mm(embbedings, kernel_norm)
        else:
            x = embbedings
            sub_kernels = torch.chunk(self.kernel, len(self.device_id), dim=1)
            temp_x = x.cuda(self.device_id[0])
            kernel_norm = l2_norm(sub_kernels[0], axis = 0).cuda(self.device_id[0])
            cos_theta = torch.mm(temp_x, kernel_norm)
            for i in range(1, len(self.device_id)):
                temp_x = x.cuda(self.device_id[i])
                kernel_norm = l2_norm(sub_kernels[i], axis = 0).cuda(self.device_id[i])
                cos_theta = torch.cat((cos_theta, torch.mm(temp_x, kernel_norm).cuda(self.device_id[0])), dim=1)

        cos_theta = cos_theta.clamp(-1, 1)  # for numerical stability
        phi = cos_theta - self.m
        label = label.view(-1, 1)  # size=(B,1)
        index = cos_theta.data * 0.0  # size=(B,Classnum)
        index.scatter_(1, label.data.view(-1, 1), 1)
        index = index.byte()
        output = cos_theta * 1.0
        output[index] = phi[index]  # only change the correct predicted output
        output *= self.s  # scale up in order to make softmax work, first introduced in normface

        return output

 def forward(self, attn_mem, n_step):
     """atten_mem: Tensor of size [num_sents, input_dim]"""
     attn_feat = torch.mm(attn_mem, self._attn_wm)
     hop_feat = torch.mm(attn_mem, self._hop_wm)
     outputs = []
     lstm_in = self._init_i.unsqueeze(0)
     lstm_states = (self._init_h.unsqueeze(1), self._init_c.unsqueeze(1))
     for _ in range(n_step):
         h, c = self._lstm_cell(lstm_in, lstm_states)
         query = h[:, -1, :]
         for _ in range(self._n_hop):
             query = PtrExtractorRL.attention(hop_feat, query,
                                             self._hop_v, self._hop_wq)
         score = PtrExtractorRL.attention_score(
             attn_feat, query, self._attn_v, self._attn_wq)
         if self.training:
             prob = F.softmax(score, dim=-1)
             out = torch.distributions.Categorical(prob)
         else:
             for o in outputs:
                 score[0, o[0, 0].item()][0] = -1e18
             out = score.max(dim=1, keepdim=True)[1]
         outputs.append(out)
         lstm_in = attn_mem[out[0, 0].item()].unsqueeze(0)
         lstm_states = (h, c)
     return outputs

    def backward(ctx, grad_output):
        matrix1, matrix2 = ctx.saved_variables
        grad_add_matrix = grad_matrix1 = grad_matrix2 = None

        if ctx.needs_input_grad[0]:
            grad_add_matrix = maybe_unexpand(grad_output, ctx.add_matrix_size)
            if ctx.alpha != 1:
                grad_add_matrix = grad_add_matrix.mul(ctx.alpha)

        if ctx.needs_input_grad[1]:
            if matrix1.stride() == (1, matrix1.size(0)):
                # column major gradient if input is column major
                grad_matrix1 = torch.mm(matrix2, grad_output.t()).t()
            else:
                grad_matrix1 = torch.mm(grad_output, matrix2.t())
            if ctx.beta != 1:
                grad_matrix1 *= ctx.beta

        if ctx.needs_input_grad[2]:
            if matrix2.stride() == (1, matrix2.size(0)):
                # column major gradient if input is column major
                grad_matrix2 = torch.mm(grad_output.t(), matrix1).t()
            else:
                grad_matrix2 = torch.mm(matrix1.t(), grad_output)
            if ctx.beta != 1:
                grad_matrix2 *= ctx.beta

        return grad_add_matrix, grad_matrix1, grad_matrix2, None, None, None

    def updateGradInput(self, input, gradOutput):
        if self.gradInput is None:
            return

        self._assertInputGradOutput(input, gradOutput)
        # compute d output / d input:
        self.gradInput[0].resize_as_(input[0]).fill_(0)
        self.gradInput[1].resize_as_(input[1]).fill_(0)

        #: first slice of weight tensor (k = 1)
        self.gradInput[0].addmm_(input[1], self.weight[0].t())
        self.gradInput[0].mul_(gradOutput.narrow(1, 0, 1).expand(self.gradInput[0].size(0),
                                                                 self.gradInput[0].size(1)))
        self.gradInput[1].addmm_(input[0], self.weight[0])
        self.gradInput[1].mul_(gradOutput.narrow(1, 0, 1).expand(self.gradInput[1].size(0),
                                                                 self.gradInput[1].size(1)))

        #: remaining slices of weight tensor
        if self.weight.size(0) > 1:
            if self.buff1 is None:
                self.buff1 = input[0].new()
            self.buff1.resize_as_(input[0])

            for k in range(1, self.weight.size(0)):
                torch.mm(input[1], self.weight[k].t(), out=self.buff1)
                self.buff1.mul_(gradOutput.narrow(1, k, 1).expand(self.gradInput[0].size(0),
                                                                  self.gradInput[0].size(1)))
                self.gradInput[0].add_(self.buff1)

                torch.mm(input[0], self.weight[k], out=self.buff2)
                self.buff2.mul_(gradOutput.narrow(1, k, 1).expand(self.gradInput[1].size(0),
                                                                  self.gradInput[1].size(1)))
                self.gradInput[1].add_(self.buff2)

        return self.gradInput

    def backward(ctx, grad_output):
        input1, input2, weight, bias = ctx.saved_variables
        grad_input1 = grad_input2 = grad_weight = grad_bias = None

        buff = Variable(input1.data.new())

        if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]:
            grad_input1 = torch.mm(input2, weight[0].t())
            grad_input1 = grad_input1.mul(grad_output.narrow(1, 0, 1).expand(grad_input1.size()))
            grad_input2 = torch.mm(input1, weight[0])
            grad_input2 = grad_input2.mul(grad_output.narrow(1, 0, 1).expand(grad_input2.size()))

            for k in range(1, weight.size(0)):
                buff = input2.mm(weight[k].t())
                buff = buff.mul(grad_output.narrow(1, k, 1).expand(grad_input1.size()))
                grad_input1.add_(buff)

                buff = input1.mm(weight[k])
                buff = buff.mul(grad_output.narrow(1, k, 1).expand(grad_input2.size()))
                grad_input2.add_(buff)

        grad_weight = Variable(weight.data.new(weight.size()))
        if ctx.needs_input_grad[2]:
            # accumulate parameter gradients:
            for k in range(weight.size(0)):
                buff = input1.mul(grad_output.narrow(1, k, 1).expand_as(input1))
                grad_weight[k] = torch.mm(buff.t(), input2)

        if bias is not None and ctx.needs_input_grad[3]:
            grad_bias = grad_output.sum(0, keepdim=False)

        return grad_input1, grad_input2, grad_weight, grad_bias

def l2l_validate(model, cluster_center, n_epoch=100):
    val_accuracy = []
    for epoch in range(n_epoch):
        data_l = generate_data_l(cluster_center)
        data_n = generate_data_n(cluster_center, model.n_class_n)
        x_l, y_l = Variable(torch.from_numpy(data_l[0])).float(), Variable(
            torch.from_numpy(data_l[1]))
        x_n, y_n = Variable(torch.from_numpy(data_n[0])).float(), Variable(
            torch.from_numpy(data_n[1]))
        pred_ll, pred_nl, w, b = model(x_l, x_n)
        M = Variable(torch.zeros(model.n_class_n, model.n_dim))
        B = Variable(torch.zeros(model.n_class_n))
        for k in range(model.n_class_n):
            M[k] = torch.cat((w[:, 0][y_n == model.n_class_l + k].view(-1, 1),
                              w[:, 1][y_n == model.n_class_l + k].view(-1, 1)), 1).mean(0)
            B[k] = b[y_n == model.n_class_l + k].mean()
        pred_ln = torch.mm(x_l, M.t()) + B.view(1, -1).expand_as(torch.mm(x_l, M.t()))
        pred_nn = torch.mm(x_n, M.t()) + B.view(1, -1).expand_as(torch.mm(x_n, M.t()))
        pred = torch.cat((torch.cat((pred_ll, pred_nl)), torch.cat((pred_ln, pred_nn))), 1)
        pred = pred.data.max(1)[1]
        y = torch.cat((y_l, y_n))
        accuracy = pred.eq(y.data).cpu().sum() * 1.0 / y.size()[0]
        # print('accuracy: %.2f' % accuracy)
        val_accuracy.append(accuracy)
        acc_l = pred.eq(y.data).cpu()[0:100].sum() * 1.0 / 100
        acc_n = pred.eq(y.data).cpu()[100:150].sum() * 1.0 / 50
        print('accuracy: %.2f, lifelong accuracy: %.2f, new accuracy: %.2f' % (accuracy, acc_l, acc_n))

    return numpy.mean(numpy.asarray(val_accuracy))

def l2l_train(model, cluster_center, n_epoch=10000, trunc_step=10):
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    M_all = Variable(torch.zeros(model.n_class, model.n_dim))
    B_all = Variable(torch.zeros(model.n_class))
    for epoch in range(n_epoch):
        loss = 0
        M_step, B_step = [], []
        for step in range(trunc_step):
            data = generate_data(cluster_center)
            optimizer.zero_grad()
            x, y = Variable(torch.from_numpy(data[0])).float(), Variable(torch.from_numpy(data[1]))
            w, b = model(x)
            M = Variable(torch.zeros(model.n_class_n, model.n_dim))
            B = Variable(torch.zeros(model.n_class_n))
            for k in range(model.n_class_n):
                M[k] = torch.cat((w[:, 0][y == model.n_class_l + k].view(-1, 1),
                                  w[:, 1][y == model.n_class_l + k].view(-1, 1)), 1).mean(0)
                B[k] = b[y == model.n_class_l + k].mean()
            if step == 0:
                M_ = M
                B_ = B
            else:
                M_ = step / (step + 1) * M_step[-1] + 1 / (step + 1) * M
                B_ = step / (step + 1) * B_step[-1] + 1 / (step + 1) * B
            M_step.append(M_)
            B_step.append(B_)
            pred = torch.mm(x, M_.t()) + B_.view(1, -1).expand_as(torch.mm(x, M_.t()))
            loss += F.cross_entropy(pred, y)
        loss.backward()
        optimizer.step()
        print('Train Epoch: {}\tLoss: {:.6f}'.format(epoch, loss.data[0]))
    return M_all, B_all, cluster_center

    def forward(self, input, hidden, encoder_outputs, enc_padding_mask,
                context, extra_zeros, enc_batch_extend_vocab, coverage):
        """
        :param input: (B)
        :param hidden: (1, B, H), (1, B, H)
        :param encoder_outputs: (B, L, 2*H)
        :param enc_padding_mask: (B, L)
        :param context: (B, 2*H); Since beam search will use context, so we need to send context out.
        :param extra_zeros: (B, n)
        :param enc_batch_extend_vocab: (B, L)
        :param coverage: (B, L)
        :return: (B, V), ((1, B, H), (1, B, H)), (B, 2*H), (B, L), (B, 1), (B, L)
        """

        input = self.embed(input)  # B -> (B, D)
        x = self.x_context(torch.cat((context, input), 1))  # (B, 2*H), (B, D) -> (B, 2*H + D) -> (B, D)
        output, hidden = self.lstm(x.unsqueeze(1), hidden)  # (B, 1, D), ((1, B, H), (1, B, H)) -> (B, 1, H), hidden

        h_decoder, c_decoder = hidden  # (1, B, H), (1, B, H)
        hidden_hat = torch.cat((h_decoder.view(-1, self.args.hidden_dim),
                                c_decoder.view(-1, self.args.hidden_dim)), 1)  # (B, H), (B, H) -> (B, 2*H)
        context, attn_dist, coverage = self.attention(hidden_hat, encoder_outputs, enc_padding_mask, coverage)
        # (B, 2*H), (B, L), (B, L) <- (B, 2*H), (B, L, 2*H), (B, L), (B, L)

        p_gen = None
        if self.args.pointer_gen:
            p_gen_input = torch.cat((context, hidden_hat, x), 1)  # (B, 2*H), (B, 2*H), (B, D) -> (B, 2*2*H + D)
            p_gen = self.p_gen_linear(p_gen_input)  # (B, 2*2*H + D) -> (B, 1)
            p_gen = torch.sigmoid(p_gen)  # (B, 1)

        output = torch.cat((output.view(-1, self.args.hidden_dim), context), 1)  # (B, H), (B, 2*H) -> (B, 3*H)
        output = self.out_linear(output)  # (B, 3*H) -> (B, H)
        # output = F.relu(output)

        ## map (B, H) -> (B, V)
        # output = self.out2(output)  # (B, H) -> (B, V); change to below matrix multiply
        output_pos = self.hidden2dim_pos(output)  # (B, H) -> (B, D)
        output_neg = self.hidden2dim_neg(output)  # (B, H) -> (B, D)
        output_pos = F.relu(torch.mm(output_pos, self.embed.weight.t()))  # (B, D) * (D, V) -> (B, V)
        output_neg = F.relu(torch.mm(output_neg, self.embed.weight.t()))  # (B, D) * (D, V) -> (B, V)
        output = output_pos - output_neg  # (B, V)

        ## change output to vocab_dist
        vocab_dist = F.softmax(output, dim=1)  # (B, V)

        if self.args.pointer_gen:
            vocab_dist_ = p_gen * vocab_dist  # (B, 1) * (B, V) -> (B, V)
            attn_dist_ = (1 - p_gen) * attn_dist  # (B, 1) * (B, L) -> (B, L)

            if extra_zeros is not None:
                vocab_dist_ = torch.cat([vocab_dist_, extra_zeros], 1)  # (B, V), (B, n) -> (B, V + n)

            final_dist = vocab_dist_.scatter_add(1, enc_batch_extend_vocab, attn_dist_)  # (B, V) -> (B, V + n)
        else:
            final_dist = vocab_dist  # (B, V)

        return final_dist, hidden, context, attn_dist, p_gen, coverage  # (B, V), ((1, B, H), (1, B, H)), (B, 2*H), (B, L), (B, 1), (B, L)

def addDecovRegularizer(loss, regParam, activations) :  
    for i in range( len(activations) ) :
        x =   activations[i]  
        batch_size = x.shape[0] #print("x.shape is: " , x.shape) # 2048, 100
        h_centered = x - torch.mean(x, dim=0, keepdim=True) # mean center activations
        covariance = torch.mm( h_centered.t(),  h_centered) # get small x small covariance matrix
        n = covariance.shape[0]
        covariance[np.diag_indices(n)] = 0 # zero out the diagonals of the covariance matrix (as we don't want to penalize the neurons against themselves) # alternative: t[torch.eye(n).byte()] = 5
        covariance /= batch_size # normalize by the length of the minibatch
        cost = ( 0.5 * regParam) * torch.sum( torch.mm(covariance, covariance) )
        loss += cost 

    def train(self, x):
        self.model.train()
        
        o = self.model(x)
        loss = torch.mean(torch.pow(torch.mm(o, self.B.t()) - x, 2))

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        U, _, V = torch.svd(torch.mm(x.t().data, o.data))
        self.B = torch.autograd.Variable(torch.mm(U, V.t()))
        
        return loss.data.cpu().numpy()

 def merge(tbl):
     inp = scn.InputBatch(2, spatial_size)
     center = spatial_size.float().view(1, 2) / 2
     p = torch.LongTensor(2)
     v = torch.FloatTensor([1, 0, 0])
     for char in tbl['input']:
         inp.addSample()
         m = torch.eye(2)
         r = random.randint(1, 3)
         alpha = random.uniform(-0.2, 0.2)
         if alpha == 1:
             m[0][1] = alpha
         elif alpha == 2:
             m[1][0] = alpha
         else:
             m = torch.mm(m, torch.FloatTensor(
                 [[math.cos(alpha), math.sin(alpha)],
                  [-math.sin(alpha), math.cos(alpha)]]))
         c = center + torch.FloatTensor(1, 2).uniform_(-8, 8)
         for stroke in char:
             stroke = stroke.float() / 255 - 0.5
             stroke = c.expand_as(stroke) + \
                 torch.mm(stroke, m * (Scale - 0.01))
             ###############################################################
             # To avoid GIL problems use a helper function:
             scn.dim_fn(
                 2,
                 'drawCurve')(
                 inp.metadata.ffi,
                 inp.features,
                 stroke)
             ###############################################################
             # Above is equivalent to :
             # x1,x2,y1,y2,l=0,stroke[0][0],0,stroke[0][1],0
             # for i in range(1,stroke.size(0)):
             #     x1=x2
             #     y1=y2
             #     x2=stroke[i][0]
             #     y2=stroke[i][1]
             #     l=1e-10+((x2-x1)**2+(y2-y1)**2)**0.5
             #     v[1]=(x2-x1)/l
             #     v[2]=(y2-y1)/l
             #     l=max(x2-x1,y2-y1,x1-x2,y1-y2,0.9)
             #     for j in numpy.arange(0,1,1/l):
             #         p[0]=math.floor(x1*j+x2*(1-j))
             #         p[1]=math.floor(y1*j+y2*(1-j))
             #         inp.setLocation(p,v,False)
             ###############################################################
     inp.precomputeMetadata(precomputeStride)
     return {'input': inp, 'target': torch.LongTensor(tbl['target']) - 1}

def anomalyScore(args, model, dataset, mean, cov, channel_idx=0, score_predictor=None):
    predictions = []
    rearranged = []
    errors = []
    hiddens = []
    predicted_scores = []
    with torch.no_grad():
        # Turn on evaluation mode which disables dropout.
        model.eval()
        pasthidden = model.init_hidden(1)
        for t in range(len(dataset)):
            out, hidden = model.forward(dataset[t].unsqueeze(0), pasthidden)
            predictions.append([])
            rearranged.append([])
            errors.append([])
            hiddens.append(model.extract_hidden(hidden))
            if score_predictor is not None:
                predicted_scores.append(score_predictor.predict(model.extract_hidden(hidden).numpy()))

            predictions[t].append(out.data.cpu()[0][0][channel_idx])
            pasthidden = model.repackage_hidden(hidden)
            for prediction_step in range(1, args.prediction_window_size):
                out, hidden = model.forward(out, hidden)
                predictions[t].append(out.data.cpu()[0][0][channel_idx])

            if t >= args.prediction_window_size:
                for step in range(args.prediction_window_size):
                    rearranged[t].append(
                        predictions[step + t - args.prediction_window_size][args.prediction_window_size - 1 - step])
                rearranged[t] =torch.FloatTensor(rearranged[t]).to(args.device).unsqueeze(0)
                errors[t] = rearranged[t] - dataset[t][0][channel_idx]
            else:
                rearranged[t] = torch.zeros(1,args.prediction_window_size).to(args.device)
                errors[t] = torch.zeros(1, args.prediction_window_size).to(args.device)

    predicted_scores = np.array(predicted_scores)
    scores = []
    for error in errors:
        mult1 = error-mean.unsqueeze(0) # [ 1 * prediction_window_size ]
        mult2 = torch.inverse(cov) # [ prediction_window_size * prediction_window_size ]
        mult3 = mult1.t() # [ prediction_window_size * 1 ]
        score = torch.mm(mult1,torch.mm(mult2,mult3))
        scores.append(score[0][0])

    scores = torch.stack(scores)
    rearranged = torch.cat(rearranged,dim=0)
    errors = torch.cat(errors,dim=0)

    return scores, rearranged, errors, hiddens, predicted_scores

 def forward(self, input_n, hidden, phi, nh):
     self.batch_size = input_n.size()[0]
     hidden = torch.cat((hidden, input_n), 2)
     # Aggregate reresentations
     h_conv = torch.div(torch.bmm(phi, hidden), nh)
     hidden = hidden.view(-1, self.hidden_size + self.input_size)
     h_conv = h_conv.view(-1, self.hidden_size + self.input_size)
     # h_conv has shape (batch_size, n, hidden_size + input_size)
     m1 = (torch.mm(hidden, self.W1)
           .view(self.batch_size, -1, self.hidden_size))
     m2 = (torch.mm(h_conv, self.W2)
           .view(self.batch_size, -1, self.hidden_size))
     m3 = self.b.unsqueeze(0).unsqueeze(1).expand_as(m2)
     hidden = torch.sigmoid(m1 + m2 + m3)
     return hidden

def count_accuracy(X, true_counts, air, batch_size):
    assert X.size(0) == true_counts.size(0), 'Size mismatch.'
    assert X.size(0) % batch_size == 0, 'Input size must be multiple of batch_size.'
    counts = torch.LongTensor(3, 4).zero_()
    error_latents = []
    error_indicators = []

    def count_vec_to_mat(vec, max_index):
        out = torch.LongTensor(vec.size(0), max_index + 1).zero_()
        out.scatter_(1, vec.type(torch.LongTensor).view(vec.size(0), 1), 1)
        return out

    for i in range(X.size(0) // batch_size):
        X_batch = X[i * batch_size:(i + 1) * batch_size]
        true_counts_batch = true_counts[i * batch_size:(i + 1) * batch_size]
        z_where, z_pres = air.guide(X_batch, batch_size)
        inferred_counts = sum(z.cpu() for z in z_pres).squeeze().data
        true_counts_m = count_vec_to_mat(true_counts_batch, 2)
        inferred_counts_m = count_vec_to_mat(inferred_counts, 3)
        counts += torch.mm(true_counts_m.t(), inferred_counts_m)
        error_ind = 1 - (true_counts_batch == inferred_counts)
        error_ix = error_ind.nonzero().squeeze()
        error_latents.append(latents_to_tensor((z_where, z_pres)).index_select(0, error_ix))
        error_indicators.append(error_ind)

    acc = counts.diag().sum().float() / X.size(0)
    error_indices = torch.cat(error_indicators).nonzero().squeeze()
    if X.is_cuda:
        error_indices = error_indices.cuda()
    return acc, counts, torch.cat(error_latents), error_indices

        def test_shape(di, dj, dk):
            x = self._gen_sparse(2, 20, [di, dj])[0]
            y = self.randn(dj, dk)

            res = torch.hsmm(x, y)
            expected = torch.mm(x.to_dense(), y)
            self.assertEqual(res.to_dense(), expected)

 def _step(self, tok, states, attention):
     prev_states, prev_out = states
     lstm_in = torch.cat(
         [self._embedding(tok).squeeze(1), prev_out],
         dim=1
     )
     states = self._lstm(lstm_in, prev_states)
     lstm_out = states[0][-1]
     query = torch.mm(lstm_out, self._attn_w)
     attention, attn_mask = attention
     context, score = step_attention(
         query, attention, attention, attn_mask)
     dec_out = self._projection(torch.cat([lstm_out, context], dim=1))
     states = (states, dec_out)
     logit = torch.mm(dec_out, self._embedding.weight.t())
     return logit, states, score

    def _step(self, tok, states, attention):
        prev_states, prev_out = states
        lstm_in = torch.cat(
            [self._embedding(tok).squeeze(1), prev_out],
            dim=1
        )
        states = self._lstm(lstm_in, prev_states)
        lstm_out = states[0][-1]
        query = torch.mm(lstm_out, self._attn_w)
        attention, attn_mask, extend_src, extend_vsize = attention
        context, score = step_attention(
            query, attention, attention, attn_mask)
        dec_out = self._projection(torch.cat([lstm_out, context], dim=1))

        # extend generation prob to extended vocabulary
        gen_prob = self._compute_gen_prob(dec_out, extend_vsize)
        # compute the probabilty of each copying
        copy_prob = torch.sigmoid(self._copy(context, states[0][-1], lstm_in))
        # add the copy prob to existing vocab distribution
        lp = torch.log(
            ((-copy_prob + 1) * gen_prob
            ).scatter_add(
                dim=1,
                index=extend_src.expand_as(score),
                source=score * copy_prob
        ) + 1e-8)  # numerical stability for log
        return lp, (states, dec_out), score

        def test_shape(di, dj, dk):
            x = self._gen_sparse(2, 20, [di, dj])[0]
            y = self.randn(dj, dk)

            res = torch.dsmm(x, y)
            expected = torch.mm(self.safeToDense(x), y)
            self.assertEqual(res, expected)

    def forward(self, X, posterior_mean = False):
        """
        Funciton call to generate the output, every time we call it, the dynamic graph is created.
        There can be difference between forward in training and test:
            - In dropout we do not zero neurons in test
            - In Variational Inference we dont randombly sample from the posterior
        
        We create the forward pass by performing operations between the input X (Nsam_batch, Ndim)
        and the parameters of the model that we should have initialized in the __init__
        """
        
        ## We need to sample from the posterior !! 
        self.sample_posterior(posterior_mean)
        
        o1 = self.linear1(X)
#        o1 = torch.mm(X, self.W1) + self.b1
#        print ("x shape: ", X.shape, "W1 shape: ", self.W1.shape, "b1 shape: ", self.b1.shape)
#        print ("o1 shape: ", o1.shape)
#        print ("W2 shape: ", self.W2.shape, "b2 shape: ", self.b2.shape)
        
        ## Apply non-linearity
        o1 = self.cf_a.activation_func(o1)
        o1 = F.dropout(o1,p = self.cf_a.dop, training = self.training)
        o2 = torch.mm(o1, self.W2) + self.b2
#        print ("o2 shape: ", o2.shape)
        return o2