Hi, I am trying to tweak (I am not a power user, sorry) a model using forward propagation and I get the warning about argument changes for the softmax call and also I get an OOM in GPU usage at the same level, can somebody give me a hint how to introduce the dim=X argument in the softmax call? Maybe from this reason I get the OOM…Thanks
UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.
Warning line:    alpha_i = F.softmax(self.U(x_i).t())
######### Method code ##########
def forward(self, x, target):
    # get embeddings and apply dropout
    x = self.embed(x)
    x = self.embed_drop(x)
    x = x.transpose(1, 2)

    # apply convolution and nonlinearity (tanh)
    x = F.tanh(self.conv(x).transpose(1,2))
    y = []
    ss = []
    alphas = []
    for x_i in x:
        # apply attention
        
        # HERE COMES THE WARNING!!!!!!!!!!!!!!!

        alpha_i = F.softmax(self.U(x_i).t())

        # document representations are weighted sums using the attention. Can compute all at once as a matmul
        m_i = alpha_i.mm(x_i)

        # final layer classification
        y_i = self.final.mul(m_i).sum(dim=1).add(self.final_bias)

        # save attention
        alphas.append(alpha_i)
        y.append(y_i)

    r = torch.stack(y)
    # print(torch.cuda.is_available())
    torch.cuda.empty_cache()
    # print('current memory allocated before: {}'.format(torch.cuda.memory_allocated() / 1024 ** 2))
    # print('max memory allocated before: {}'.format(torch.cuda.max_memory_allocated() / 1024 ** 2))
    # print('cached memory before: {}'.format(torch.cuda.memory_cached() / 1024 ** 2))
    alpha = torch.stack(alphas)
    torch.cuda.empty_cache()
    # final sigmoid to get predictions
    yhat = F.sigmoid(r)
    if target is not None:
        loss = self.get_loss(yhat, target)
    else:
        loss = 0
    return yhat, loss, alpha