Solved by ptrblck in post #2 
              
              
                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. 
Are you using different versions of the repository?