Closed KumoLiu closed 2 months ago
--------------------------------------------------------------------------- TypeError Traceback (most recent call last) Cell In[4], line 39 36 # Model Definition 37 device = torch.device("cuda") ---> 39 model = UNet( 40 dimensions=3, 41 in_channels=1, 42 out_channels=3, 43 channels=(16, 32, 64, 128, 256), 44 strides=(2, 2, 2, 2), 45 num_res_units=2, 46 norm=Norm.BATCH, 47 dropout=0.2, 48 ).cuda() 50 # Load Model Weights 51 model.load_state_dict(torch.load(model_ckpt)) TypeError: UNet.__init__() got an unexpected keyword argument 'dimensions'