5 Commits

Author	SHA1	Message	Date
Ross Wightman	c527c37969	Optimizations for pos embed resize, merge different mask helper fns	2025-04-21 14:05:18 -07:00
Ross Wightman	ea728f67fa	Improve several typing issues for flex vit, can (almost) work with jit if we bash h,w key into an int or str	2025-04-14 11:01:56 -07:00
Ross Wightman	b4bb0f452a	Exclude embeds module and mask attn functions from tracing	2025-04-09 15:34:15 -07:00
Ross Wightman	13e0f3a4a3	Add loss scale arg, initial distributed loss scale. Maybe fix FX for the model.	2025-04-08 20:47:57 -07:00
Ross Wightman	0893f5d296	Initial NaFlex ViT model and training support	2025-04-07 21:27:10 -07:00