import torch
import numpy as np

device = "cpu"
input_array = np.zeros((4, 16, 16), dtype=np.float32)
input_tensor = torch.from_numpy(input_array).unsqueeze(0).to(device)

print("Before pad:", input_tensor.shape)

from train_cloud_removal import UNet
model = UNet(in_channels=6, out_channels=4).to(device)

if hasattr(model, 'inc') and hasattr(model.inc.double_conv[0], 'in_channels'):
    expected_channels = model.inc.double_conv[0].in_channels
elif hasattr(model, 'conv1') and hasattr(model.conv1, 'in_channels'):
    expected_channels = model.conv1.in_channels
else:
    expected_channels = list(model.parameters())[0].shape[1]

print("Expected channels:", expected_channels)

if expected_channels > input_tensor.shape[1]:
    pad_channels = expected_channels - input_tensor.shape[1]
    padding = torch.zeros(1, pad_channels, *input_tensor.shape[2:]).to(device)
    input_tensor = torch.cat([input_tensor, padding], dim=1)

print("After pad:", input_tensor.shape)

try:
    model(input_tensor)
    print("Success!")
except Exception as e:
    print("Error:", e)