I’m having issue with input/output sizes.
This is my dataloader:
train_data = torch.hstack((train_feat, train_labels))
train_loader = torch.utils.data.DataLoader(train_data, batch_size= 128, shuffle=True)

This is my dataset:
print(len(train_loader))
24079

print((train_feat.shape))
torch.Size([3082092, 13])

print((train_labels.shape))
torch.Size([3082092, 1])

This is my model:
class TDNN(nn.Module):

    def __init__(self, feat_dim=13, embedding_size=512, num_classes=51,
                 config_str='batchnorm-relu'):
        super(TDNN, self).__init__()

        self.network = nn.Sequential(OrderedDict([
            ('tdnn1', TDNNLayer(feat_dim, 512, 5, dilation=1, padding=0,
                                config_str=config_str)),
            ('tdnn2', TDNNLayer(512, 512, 3, dilation=2, padding=0,
                                config_str=config_str)),
            ('tdnn3', TDNNLayer(512, 512, 3, dilation=3, padding=0,
                                config_str=config_str)),
            ('tdnn4', DenseLayer(512, 512, config_str=config_str)),
            ('tdnn5', DenseLayer(512, 1500, config_str=config_str)),
            ('stats', StatsPool()),
            ('affine', nn.Linear(3000, embedding_size))
        ]))
        self.nonlinear = get_nonlinear(config_str, embedding_size)
        self.dense = DenseLayer(embedding_size, embedding_size, config_str=config_str)
        self.classifier = nn.Linear(embedding_size, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight.data)
                nn.init.zeros_(m.bias)

    def forward(self, x):
        x = self.network(x)
        if self.training:
            x = self.dense(self.nonlinear(x))
            x = self.classifier(x)
        return x

The input should be batches of 128. I’m probably making an error with the data here:
class IterMeter(object):
    """keeps track of total iterations"""
    def __init__(self):
        self.val = 0

    def step(self):
        self.val += 1

    def get(self):
        return self.val


def train(model, device, train_loader, criterion, optimizer, scheduler, epoch, iter_meter, experiment):
    model.train()
    data_len = len(train_loader.dataset)
    with experiment.train():
        for batch_idx, _data in enumerate(train_loader):
            features, labels = _data[:, :-1], _data[:, -1] 
            features, labels = features.to(device), labels.to(device)
            features = features.unsqueeze(dim=1)
            
            optimizer.zero_grad()

            output = model(features)  # (batch, n_class)

            loss = criterion(output, labels)
            loss.backward()

            experiment.log_metric('loss', loss.item(), step=iter_meter.get())
            experiment.log_metric('learning_rate', scheduler.get_lr(), step=iter_meter.get())

            optimizer.step()
            scheduler.step()
            iter_meter.step()
            if batch_idx % 100 == 0 or batch_idx == data_len:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(features), data_len,
                    100. * batch_idx / len(train_loader), loss.item()))

iter_meter = IterMeter()
for epoch in range(1, epochs + 1):
    train(model, device, train_loader, criterion, optimizer, scheduler, epoch, iter_meter, experiment)

If I remore the features.unsqueeze(dim=1), I get:
RuntimeError: Expected 3-dimensional input for 3-dimensional weight [512, 13, 5], but got 2-dimensional input of size [128, 13] instead