search

lavish619 / MLP-Mixer-PyTorch

An all MLP architecture for Computer Vision by Google (Paper Implementation)
MIT License
6 stars 4 forks source link

class MLP_Mixer miss use in forward causing model stored in cpu #1

Closed lclkent closed 3 years ago

lclkent commented 3 years ago

The model should be initialized in init as bellow and inferred in forward layers = [] for _ in range(self.num_mixers): layers.append(Mixer_Layer(self.num_patches, self.num_features, self.expansion_token, self.expansion_channel, self.dropout)) self.model = nn.Sequential(*layers)

lavish619 commented 3 years ago

Thanks for pointing this out, I haven't tested it on GPU initially. Now I have corrected it.