34 lines
1.0 KiB
Python
34 lines
1.0 KiB
Python
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)
|