Hey KFrank,
Thanks again for the reply.
Here’s the code snippet which you asked for.
class MLPUpdated(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(30, 32),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(32,16),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(16, 8),
            nn.ReLU(),
            nn.Linear(8, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.layers(x)

device = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device)
model = MLPUpdated().to(device)

train_data = torch.tensor([ 0.3018, -0.3064,  1.1054, -0.0751,  0.4352,  0.9034, -0.4978, -0.4954,
         0.0315, -0.0492,  0.4202, -0.2261, -0.2327,  0.4352, -0.5217,  0.1191,
        -0.0143, -0.7541,  0.4915, -1.1997, -0.2744, -0.1156, -0.0913, -0.1326,
        -0.0158,  0.4527,  0.4410,  0.3742, -0.0228,  0.0137])
print(train_data.shape)

test_data = torch.tensor([ 0.5229,  0.8923, -1.6663, -1.8197,  1.1027, -1.2481,  0.3701,  1.5603,
        -1.6727,  1.3396,  0.1646, -0.6536, -0.4564,  1.0487, -0.0249, -0.4200,
        -1.7048, -1.9665,  0.4084,  0.8154, -1.8285, -0.3640, -0.0078,  0.1991,
         0.2119, -1.6663, -1.0810, -0.4766,  0.2329, -0.0569])
print(test_data.shape)

model(train_data)
model.eval()
with torch.no_grad():
    model(test_data)
print("SUCCESS")

Output:
2021-06-26 12:06:47.330716: W tensorflow/stream_executor/platform/default/dso_loader.cc:59] Could not load dynamic library 'cudart64_101.dll'; dlerror: cudart64_101.dll not found
2021-06-26 12:06:47.340872: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.        
torch.Size([30])
torch.Size([30])
SUCCESS

This is running successfully.
But when I run the below code, it gets the error of matrices expected.
#!/usr/bin/python3

"""federated.py Contains an implementation of federated learning with ten
                workers applied to the Fashion MNIST data set for image
                classification using a slightly modified version of the LeNet5
                CNN architecture.

For the ID2223 Scalable Machine Learning course at KTH Royal Institute of
Technology"""

__author__ = "Xenia Ioannidou and Bas Straathof"


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

from sys import argv
from time import time
from argparse import ArgumentParser, Namespace
from syft.federated.floptimizer import Optims

# The Pysyft Federated Learning library
import syft as sy

from models import LeNetComplete, MLP, MLPUpdated
from plotting import generate_simple_plot

from load_data import PeopleDataset

torch.autograd.set_detect_anomaly(True)

def parse_args() -> Namespace:
    """Parses CL arguments

    Returns:
        Namespace object containing all arguments"""
    parser = ArgumentParser()

    parser.add_argument("-bs", "--batch_size", type=int, default=64)
    parser.add_argument("-tbs", "--test_batch_size", type=int, default=1000)
    parser.add_argument("-ls", "--log_steps", type=int, default=50)
    parser.add_argument("-lr", "--learning_rate", type=float, default=0.01)
    parser.add_argument("-e", "--epochs", type=int, default=25)

    return parser.parse_args(argv[1:])


def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    """Train the model

    Args:
        args: Hyper-paramters for training
        model: The model to be trained
        device: Training device (e.g. cpu/gpu)
        train_loader: Federated data loader for training
        optimizer: Training optimizer (e.g. SGD)
        epoch: Current epoch

    Returns:
        training_time: Time it took for one epoch of training
    """
    start = time()
    # data, target = next(iter(train_loader))
    # batch_idx = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        # Ensure that thet data is send to the right worker
        # print("Train: " ,data.shape)
        model.send(data.location)

        # Write the data and targets to the compute device
        data, target = data.to(device), target.to(device)

        #Call the optimizer for the worker using get_optim
        # print(data.location)
        # print(optimizer.optimizers)
        opt = optimizer.get_optim(data.location.id)

        # Reset the optimizers gradients to zero
        # optimizer.zero_grad()
        opt.zero_grad()

        # Pass the data to the model
        print("Train: ", data.shape)
        # print("Train Type: ", type(data.location))
        output = model(data)

        # Compute the loss
        # binary_loss = nn.BCELoss()
        binary_loss = torch.nn.BCEWithLogitsLoss()
        target = target.unsqueeze(1)
        target = target.float()
        loss = binary_loss(output, target)
        # loss = F.nll_loss(output, target)

        # Execute the backpropagation step
        loss.backward()

        # Apply the optimizer
        opt.step()

        # Retrieve the model (PySyft feature)
        model.get()

        if batch_idx % args.log_steps == 0:
            # Retrieve the loss (PySyft feature)
            loss = loss.get()
            print(' {} - Epoch: {} [{}/{} ({:.0f}%)] Loss: {:.6f}'.format(
                data.location.id, epoch, batch_idx * args.batch_size,
                len(train_loader) * args.batch_size,
                100. * batch_idx / len(train_loader), loss.item()))

    end = time()
    training_time = end-start

    return training_time

def test(args, model, device, test_loader, epoch):
    """Evaluate the model

    Args:
        args: Hyper-parameters for testing
        model: The model to be evaluated
        device: Training device (e.g. cpu/gpu)
        test_loader: Federated data loader for testing
        epoch: Current epoch

    Returns:
        test_loss: Test loss for current epoch
        test_acc: Test accuracy for current epoch
    """
    model.eval()
    test_loss, correct = 0.0, 0

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)

            # Feed-forward data through the model
            print("Test: ", data.shape)
            # print("Test Type: ", type(target))
            output = model(data)

            # Sum up the batch loss
            binary_loss = torch.nn.BCEWithLogitsLoss()
            target = target.unsqueeze(1)
            target = target.float()
            test_loss += binary_loss(output, target, reduction="sum").item()
            # test_loss += F.nll_loss(output, target, reduction="sum").item()

            # Retrieve the index of the maximum log-probability
            pred = output.argmax(1, keepdim=True)

            # Get the number of correct classifications
            correct += pred.eq(target.view_as(pred)).sum().item()

    # Average the test_loss
    test_loss /= len(test_loader.dataset)

    # Test accuracy
    test_acc = correct / len(test_loader.dataset)

    print("\nTest set - Epoch: {} - Loss: {:.4f}, Acc: {:.2f}%\n".format(
        epoch, test_loss, 100 * test_acc))

    return test_loss, test_acc

if __name__ == "__main__":
    args = parse_args()

    # Pysyft needs to be hooked to PyTorch to enable its features
    hook = sy.TorchHook(torch)

    # Define the workers
    alfa    = sy.VirtualWorker(hook, id="alfa")
    bravo   = sy.VirtualWorker(hook, id="bravo")

    workers = (alfa, bravo)
    workers_id = (alfa.id, bravo.id)

    # Check if a GPU is available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    device = torch.device(device)
    kwargs = {'num_workers': 1, 'pin_memory': True} if device=="cuda" else {}


    # Federated learning needs a special train loader to distribute the data
    # over the various workers
    print("Loading data...")
    train_file = '../dataset/creditcard_train.csv'
    train_ds = PeopleDataset(train_file)

    train_loader = sy.FederatedDataLoader(train_ds.federate(workers),
        batch_size=args.batch_size, shuffle=True, **kwargs)


    test_file = '../dataset/creditcard_test.csv'
    test_ds = PeopleDataset(test_file)
    
    test_loader = sy.FederatedDataLoader(test_ds.federate(workers),
        batch_size=args.batch_size, shuffle=True, **kwargs)


    print("Data loaded...")

    # Instantiate the CNN model
    model = MLPUpdated().to(device)

    # Load the optimizer
    optimizer = Optims(workers_id, optim=optim.Adam(params=model.parameters(),lr=0.1))

    # Train and evaluate for a number of epochs
    total_train_time, test_losses, test_accs = 0.0, [], []
    for epoch in range(1, args.epochs + 1):
        train_time = train(args, model, device, train_loader, optimizer, epoch)
        total_train_time += train_time

        test_loss, test_acc =  test(args, model, device, test_loader, epoch)
        test_losses.append(test_loss)
        test_accs.append(test_acc)



    print("Total training time: {:.2f}s".format(total_train_time))
    print("Final test accuracy: {:.4f}".format(test_acc))
    print("Final test loss: {:.4f}".format(test_loss))



This gives the following output:
Train:  torch.Size([1024, 30])
Train:  torch.Size([1024, 30])
Train:  torch.Size([1024, 30])
 alfa - Epoch: 1 [102400/257024 (40%)] Loss: 0.693148
Train:  torch.Size([1024, 30])
Train:  torch.Size([1024, 30])
Train:  torch.Size([1024, 30])
.
.
.
Train:  torch.Size([1024, 30])
Train:  torch.Size([1024, 30])
 bravo - Epoch: 1 [256000/257024 (100%)] Loss: 0.693148
Train:  torch.Size([163, 30])
Test:  torch.Size([1024, 30])
Traceback (most recent call last):
.
.
.
Traceback (most recent call last):
  File "federated_learning.py", line 284, in <module>
    test_loss, test_acc =  test(args, model, device, test_loader, epoch)
  File "federated_learning.py", line 146, in test
    output = model(data)
  File "C:\Users\het\AppData\Local\Programs\Python\Python38\lib\site-packages\torch\nn\modules\module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "\Federated Learning\Split-Learning-and-Federated-Learning\src\models.py", line 206, in forward
    return self.layers(x)
.
.
.
RuntimeError: matrices expected, got 1D, 2D tensors at C:\w\1\s\windows\pytorch\aten\src\TH/generic/THTensorMath.cpp:131