What are the inputs of VarNet

[osbm]

Hello. Sorry if this looks simple but,
I am trying to use fastmri prostate data with the VarNet model. It is very hard for me to undersand what type of input does the varnet expects. I tried many things but i could not get this model to work.

from fastmri.models import VarNet
import torch

model = VarNet()

masked_kspace = torch.randn(1, 20, 256, 256)
mask = torch.randn(1, 20, 256, 256)

print(model(masked_kspace, mask).shape)

But it says

--> 186         squeezed_mask = mask[:, 0, 0, :, 0].to(torch.int8)
    187         cent = squeezed_mask.shape[1] // 2
    188         # running argmin returns the first non-zero

IndexError: too many indices for tensor of dimension 4

How does the mask should have 5 dimensions? what are those dimensions.
Normally i would try to understand this by looking at the data processing step but it looks obfuscated to me :) .
I just need to pass 1 slice of an prostate image through the network. I can easily turn that to training loop that fits my needs.

[mmuckley]

Hello @osbm For prostate data it may be best to post an issue directly on their dedicated repository https://github.com/cai2r/fastMRI_prostate.