I assume you are using @rwightman’s implementation of this model?
If so, note that the flatten argument seems to have been added 13 days ago here 2.
Are you using different versions of the repository?