Strange results from pytorch model

[linhanwang]

Thank you for your great work. Translating the model to Torch has made my life much easier. However, when I tried to use it, I got some strange output.

Here is my code:

import cosem_models
import numpy as np
import tifffile as tif
import torch
from einops import rearrange

model = cosem_models.load_model(
    "./v0003.2-pytorch/cosem_models/cosem_models/setup04/1820500"
)

print("min input shape: ", model.min_input_shape)
print("input size step: ", model.input_size_step)

img = tif.imread(
    "/home/linhan/data/dummy.tif"
)

input = img[:216, :216, :216] / np.float32(255.0)
input = 2.0 * input - 1.0
print(input.shape, input.dtype, input.min(), input.max())

# print(model)
model = model.to("cuda")
model.eval()
with torch.no_grad():
    input = torch.tensor(input, device="cuda")
    input = rearrange(input, "h w d -> 1 1 h w d")
    output = model.prediction_head(model.unet(input))
    print(output.size(), output.max(), output.min())

    output = output.cpu().numpy()

    pred = np.argmax(output, axis=1)

    print(pred.shape, pred.dtype, pred.max(), pred.min())

The output is:

min input shape:  [216. 216. 216.]
input size step:  [18 18 18]
(216, 216, 216) float32 -1.0 1.0
torch.Size([1, 14, 68, 68, 68]) tensor(2.1361e+15, device='cuda:0') tensor(-2.6005e+15, device='cuda:0')
(1, 68, 68, 68) int64 0 0

The output values are extremely large, and the predicted classes are all zero. I read single_block_inference.py and ensured the input normalization to [-1, 1]. I also checked the data type of my input, which is correct.

Do you have any idea why I am getting these strange results? Thank you so much!