def __init__(self, input_nc=512, output_dim=6):
    super(FeatureRegression, self).__init__()
    self.conv = nn.Sequential(nn.Conv2d(input_nc, 512, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.Conv2d(512, 256, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.Conv2d(256, 128, kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Conv2d(128, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True))
    self.linear = nn.Linear(((64 * 4) * 3), output_dim)
    self.tanh = nn.Tanh()
    'self.conv.to(device)\n        self.linear.to(device)\n        self.tanh.to(device)'
def __init__(self, input_nc=512, output_dim=6):
    super(FeatureRegression, self).__init__()
    self.conv = nn.Sequential(nn.Conv2d(input_nc, 512, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.Conv2d(512, 256, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.Conv2d(256, 128, kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Conv2d(128, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True))
    self.linear = nn.Linear(((64 * 4) * 3), output_dim)
    self.tanh = nn.Tanh()
    'self.conv.to(device)\n        self.linear.to(device)\n        self.tanh.to(device)'