code: def __init__(self, in_channels=1, n_class=3):
      filters = [32, 64, 128, 256, 320]
      inp, out = filters[:-1], filters[1:]
      self.input_block = DownsampleBlock(
          in_channels, filters[0], stride=1
      )
      self.downsample = [
          DownsampleBlock(i, o) for i, o in zip(inp, out)
      ]
      self.bottleneck = DownsampleBlock(filters[-1], filters[-1])
      self.upsample = [UpsampleBlock(filters[-1], filters[-1])] + [
          UpsampleBlock(i, o) for i, o in zip(out[::-1], inp[::-1])
      ]
      self.output = {
          "conv": nn.Conv2d(filters[0], n_class, kernel_size=(1, 1, 1))
      }
  
file_path: tinygrad/extra/models/unet3d.py
imports: from pathlib import Path
  import torch
  from tinygrad import nn
  from tinygrad.tensor import Tensor
  from tinygrad.helpers import fetch, get_child
label_desc: Use image node_id 1 to create a new UNet3D object  with example: obj = UNet3D(in_channels, n_class)