import os
import sys

import torch
import torch.nn as nn


class DepthEstimation(nn.Module):
    def __init__(self, device='cuda'):
        super().__init__()

        self.layer_1 = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=3, device=device, padding='same')  

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.layer_1(x)
        
        return x