Hello Kamil,
Please find the complete code given below:
import torch
import torch.nn as nn
import torch.nn.functional as F
class TransformationNet(nn.Module):
def __init__(self, input_dim, output_dim):
    super(TransformationNet, self).__init__()
    self.output_dim = output_dim

    self.conv_1 = nn.Conv1d(input_dim, 64, 1)
    self.conv_2 = nn.Conv1d(64, 128, 1)
    self.conv_3 = nn.Conv1d(128, 1024, 1)

    self.bn_1 = nn.BatchNorm1d(64)
    self.bn_2 = nn.BatchNorm1d(128)
    self.bn_3 = nn.BatchNorm1d(1024)
    self.bn_4 = nn.BatchNorm1d(512)
    self.bn_5 = nn.BatchNorm1d(256)

    self.fc_1 = nn.Linear(1024, 512)
    self.fc_2 = nn.Linear(512, 256)
    self.fc_3 = nn.Linear(256, self.output_dim*self.output_dim)

def forward(self, x):
    num_points = x.shape[1]
    x = x.transpose(2, 1)
    x = F.relu(self.bn_1(self.conv_1(x)))
    x = F.relu(self.bn_2(self.conv_2(x)))
    x = F.relu(self.bn_3(self.conv_3(x)))

    x = nn.MaxPool1d(num_points)(x)
    x = x.view(-1, 1024)

    x = F.relu(self.bn_4(self.fc_1(x)))
    x = F.relu(self.bn_5(self.fc_2(x)))
    x = self.fc_3(x)

    identity_matrix = torch.eye(self.output_dim)
    if torch.cuda.is_available():
        identity_matrix = identity_matrix.cuda()
    x = x.view(-1, self.output_dim, self.output_dim) + identity_matrix
    return x

class BasePointNet(nn.Module):
def __init__(self, point_dimension, return_local_features=False):
    super(BasePointNet, self).__init__()
    self.return_local_features = return_local_features
    self.input_transform = TransformationNet(input_dim=point_dimension, output_dim=point_dimension)
    self.feature_transform = TransformationNet(input_dim=64, output_dim=64)

    self.conv_1 = nn.Conv1d(point_dimension, 64, 1)
    self.conv_2 = nn.Conv1d(64, 64, 1)
    self.conv_3 = nn.Conv1d(64, 64, 1)
    self.conv_4 = nn.Conv1d(64, 128, 1)
    self.conv_5 = nn.Conv1d(128, 1024, 1)

    self.bn_1 = nn.BatchNorm1d(64)
    self.bn_2 = nn.BatchNorm1d(64)
    self.bn_3 = nn.BatchNorm1d(64)
    self.bn_4 = nn.BatchNorm1d(128)
    self.bn_5 = nn.BatchNorm1d(1024)

def forward(self, x):
    num_points = x.shape[1]

    input_transform = self.input_transform(x)

    x = torch.bmm(x, input_transform)
    x = x.transpose(2, 1)
    x = F.relu(self.bn_1(self.conv_1(x)))
    x = F.relu(self.bn_2(self.conv_2(x)))
    x = x.transpose(2, 1)

    feature_transform = self.feature_transform(x)

    x = torch.bmm(x, feature_transform)
    local_point_features = x

    x = x.transpose(2, 1)
    x = F.relu(self.bn_3(self.conv_3(x)))
    x = F.relu(self.bn_4(self.conv_4(x)))
    x = F.relu(self.bn_5(self.conv_5(x)))
    x = nn.MaxPool1d(num_points)(x)
    x = x.view(-1, 1024)

    if self.return_local_features:
        x = x.view(-1, 1024, 1).repeat(1, 1, num_points)
        return torch.cat([x.transpose(2, 1), local_point_features], 2), feature_transform
    else:
        return x, feature_transform

class ClassificationPointNet(nn.Module):
def __init__(self, num_classes, dropout=0.3, point_dimension=3):
    super(ClassificationPointNet, self).__init__()
    self.base_pointnet = BasePointNet(return_local_features=False, point_dimension=point_dimension)

    self.fc_1 = nn.Linear(1024, 512)
    self.fc_2 = nn.Linear(512, 256)
    self.fc_3 = nn.Linear(256, num_classes)

    self.bn_1 = nn.BatchNorm1d(512)
    self.bn_2 = nn.BatchNorm1d(256)

    self.dropout_1 = nn.Dropout(dropout)

def forward(self, x):
    x, feature_transform = self.base_pointnet(x)

    x = F.relu(self.bn_1(self.fc_1(x)))
    x = F.relu(self.bn_2(self.fc_2(x)))
    x = self.dropout_1(x)

    return F.log_softmax(self.fc_3(x), dim=1), feature_transform

Best Regards,
Vijay