Thanks! Seems like the person who made the Keras version decided to do 7 classes instead of the original 8 classes that were used in the paper. Also other thing, the paper mentioned to do xavier initialization and I read around to see how to apply it on PyTorch and it says to either do it on every individual layer or apply directly on the linear transformation. So I tried doing it on the linear transformation and got the following error:
AttributeError: 'Linear' object has no attribute 'ndimension'

This is what my code looks like so far:
import torch.nn as nn 
import torch.utils.model_zoo as model_zoo
from torch.autograd import Variable
import torch.nn.functional as F


class CNet(nn.Module):
    def __init__(self, num_classes=7):
        # Input x is (128, 128, 1
        super(HeartNet, self).__init__()

        self.features = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
            nn.ELU(inplace=True),
            nn.BatchNorm2d(64, eps=0.001),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.ELU(inplace=True),
            nn.BatchNorm2d(64, eps=0.001),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ELU(inplace=True),
            nn.BatchNorm2d(128, eps=0.001),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.ELU(inplace=True),
            nn.BatchNorm2d(128, eps=0.001),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.ELU(inplace=True),
            nn.BatchNorm2d(256, eps=0.001),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.ELU(inplace=True),
            nn.BatchNorm2d(256, eps=0.001),
            nn.MaxPool2d(kernel_size=2, stride=2)
            )

        self.classifier = nn.Sequential(
            nn.Dropout(0.5),
            nn.init.xavier_uniform(nn.Linear(16*16*256, 2048)),
            nn.ELU(inplace=True),
            nn.BatchNorm2d(2048, eps=0.001),
            nn.init.xavier_uniform(nn.Linear(2048, num_classes))
            )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), 16 * 16 * 256)
        x = self.classifier(x)
        return x

# Testing
model = CNet()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)