18 lines
383 B
Python
18 lines
383 B
Python
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
|