Repository: LAION-AI/CLAP (language: python)

Issue:
Optimize data loader. Reduced training time to half. Need to clean th…
…e comments later.


Discussion / hints:

Fix patch:
diff --git a/class_labels/FSD50k_class_labels_indices.json b/class_labels/FSD50k_class_labels_indices.json
index 0f6780b..23d98f2 100644
--- a/class_labels/FSD50k_class_labels_indices.json
+++ b/class_labels/FSD50k_class_labels_indices.json
@@ -1,1 +1,1 @@
-{"Speech": 0, "Male speech, man speaking": 1, "Female speech, woman speaking": 2, "Child speech, kid speaking": 3, "Conversation": 4, "Narration, monologue": 5, "Babbling": 6, "Speech synthesizer": 7, "Shout": 8, "Bellow": 9, "Whoop": 10, "Yell": 11, "Battle cry": 12, "Children shouting": 13, "Screaming": 14, "Whispering": 15, "Laughter": 16, "Baby laughter": 17, "Giggle": 18, "Snicker": 19, "Belly laugh": 20, "Chuckle, chortle": 21, "Crying, sobbing": 22, "Baby cry, infant cry": 23, "Whimper": 24, "Wail, moan": 25, "Sigh": 26, "Singing": 27, "Choir": 28, "Yodeling": 29, "Chant": 30, "Mantra": 31, "Male singing": 32, "Female singing": 33, "Child singing": 34, "Synthetic singing": 35, "Rapping": 36, "Humming": 37, "Groan": 38, "Grunt": 39, "Whistling": 40, "Breathing": 41, "Wheeze": 42, "Snoring": 43, "Gasp": 44, "Pant": 45, "Snort": 46, "Cough": 47, "Throat clearing": 48, "Sneeze": 49, "Sniff": 50, "Run": 51, "Shuffle": 52, "Walk, footsteps": 53, "Chewing, mastication": 54, "Biting": 55, "Gargling": 56, "Stomach rumble": 57, "Burping, eructation": 58, "Hiccup": 59, "Fart": 60, "Hands": 61, "Finger snapping": 62, "Clapping": 63, "Heart sounds, heartbeat": 64, "Heart murmur": 65, "Cheering": 66, "Applause": 67, "Chatter": 68, "Crowd": 69, "Hubbub, speech noise, speech babble": 70, "Children playing": 71, "Animal": 72, "Domestic animals, pets": 73, "Dog": 74, "Bark": 75, "Yip": 76, "Howl": 77, "Bow-wow": 78, "Growling": 79, "Whimper (dog)": 80, "Cat": 81, "Purr": 82, "Meow": 83, "Hiss": 84, "Caterwaul": 85, "Livestock, farm animals, working animals": 86, "Horse": 87, "Clip-clop": 88, "Neigh, whinny": 89, "Cattle, bovinae": 90, "Moo": 91, "Cowbell": 92, "Pig": 93, "Oink": 94, "Goat": 95, "Bleat": 96, "Sheep": 97, "Fowl": 98, "Chicken, rooster": 99, "Cluck": 100, "Crowing, cock-a-doodle-doo": 101, "Turkey": 102, "Gobble": 103, "Duck": 104, "Quack": 105, "Goose": 106, "Honk": 107, "Wild animals": 108, "Roaring cats (lions, tigers)": 109, "Roar": 110, "Bird": 111, "Bird vocalization, bird call, bird song": 112, "Chirp, tweet": 113, "Squawk": 114, "Pigeon, dove": 115, "Coo": 116, "Crow": 117, "Caw": 118, "Owl": 119, "Hoot": 120, "Bird flight, flapping wings": 121, "Canidae, dogs, wolves": 122, "Rodents, rats, mice": 123, "Mouse": 124, "Patter": 125, "Insect": 126, "Cricket": 127, "Mosquito": 128, "Fly, housefly": 129, "Buzz": 130, "Bee, wasp, etc.": 131, "Frog": 132, "Croak": 133, "Snake": 134, "Rattle": 135, "Whale vocalization": 136, "Music": 137, "Musical instrument": 138, "Plucked string instrument": 139, "Guitar": 140, "Electric guitar": 141, "Bass guitar": 142, "Acoustic guitar": 143, "Steel guitar, slide guitar": 144, "Tapping (guitar technique)": 145, "Strum": 146, "Banjo": 147, "Sitar": 148, "Mandolin": 149, "Zither": 150, "Ukulele": 151, "Keyboard (musical)": 152, "Piano": 153, "Electric piano": 154, "Organ": 155, "Electronic organ": 156, "Hammond organ": 157, "Synthesizer": 158, "Sampler": 159, "Harpsichord": 160, "Percussion": 161, "Drum kit": 162, "Drum machine": 163, "Drum": 164, "Snare drum": 165, "Rimshot": 166, "Drum roll": 167, "Bass drum": 168, "Timpani": 169, "Tabla": 170, "Cymbal": 171, "Hi-hat": 172, "Wood block": 173, "Tambourine": 174, "Rattle (instrument)": 175, "Maraca": 176, "Gong": 177, "Tubular bells": 178, "Mallet percussion": 179, "Marimba, xylophone": 180, "Glockenspiel": 181, "Vibraphone": 182, "Steelpan": 183, "Orchestra": 184, "Brass instrument": 185, "French horn": 186, "Trumpet": 187, "Trombone": 188, "Bowed string instrument": 189, "String section": 190, "Violin, fiddle": 191, "Pizzicato": 192, "Cello": 193, "Double bass": 194, "Wind instrument, woodwind instrument": 195, "Flute": 196, "Saxophone": 197, "Clarinet": 198, "Harp": 199, "Bell": 200, "Church bell": 201, "Jingle bell": 202, "Bicycle bell": 203, "Tuning fork": 204, "Chime": 205, "Wind chime": 206, "Change ringing (campanology)": 207, "Harmonica": 208, "Accordion": 209, "Bagpipes": 210, "Didgeridoo": 211, "Shofar": 212, "Theremin": 213, "Singing bowl": 214, "Scratching (performance technique)": 215, "Pop music": 216, "Hip hop music": 217, "Beatboxing": 218, "Rock music": 219, "Heavy metal": 220, "Punk rock": 221, "Grunge": 222, "Progressive rock": 223, "Rock and roll": 224, "Psychedelic rock": 225, "Rhythm and blues": 226, "Soul music": 227, "Reggae": 228, "Country": 229, "Swing music": 230, "Bluegrass": 231, "Funk": 232, "Folk music": 233, "Middle Eastern music": 234, "Jazz": 235, "Disco": 236, "Classical music": 237, "Opera": 238, "Electronic music": 239, "House music": 240, "Techno": 241, "Dubstep": 242, "Drum and bass": 243, "Electronica": 244, "Electronic dance music": 245, "Ambient music": 246, "Trance music": 247, "Music of Latin America": 248, "Salsa music": 249, "Flamenco": 250, "Blues": 251, "Music for children": 252, "New-age music": 253, "Vocal music": 254, "A capella": 255, "Music of Africa": 256, "Afrobeat": 257, "Christian music": 258, "Gospel music": 259, "Music of Asia": 260, "Carnatic music": 261, "Music of Bollywood": 262, "Ska": 263, "Traditional music": 264, "Independent music": 265, "Song": 266, "Background music": 267, "Theme music": 268, "Jingle (music)": 269, "Soundtrack music": 270, "Lullaby": 271, "Video game music": 272, "Christmas music": 273, "Dance music": 274, "Wedding music": 275, "Happy music": 276, "Funny music": 277, "Sad music": 278, "Tender music": 279, "Exciting music": 280, "Angry music": 281, "Scary music": 282, "Wind": 283, "Rustling leaves": 284, "Wind noise (microphone)": 285, "Thunderstorm": 286, "Thunder": 287, "Water": 288, "Rain": 289, "Raindrop": 290, "Rain on surface": 291, "Stream": 292, "Waterfall": 293, "Ocean": 294, "Waves, surf": 295, "Steam": 296, "Gurgling": 297, "Fire": 298, "Crackle": 299, "Vehicle": 300, "Boat, Water vehicle": 301, "Sailboat, sailing ship": 302, "Rowboat, canoe, kayak": 303, "Motorboat, speedboat": 304, "Ship": 305, "Motor vehicle (road)": 306, "Car": 307, "Vehicle horn, car horn, honking": 308, "Toot": 309, "Car alarm": 310, "Power windows, electric windows": 311, "Skidding": 312, "Tire squeal": 313, "Car passing by": 314, "Race car, auto racing": 315, "Truck": 316, "Air brake": 317, "Air horn, truck horn": 318, "Reversing beeps": 319, "Ice cream truck, ice cream van": 320, "Bus": 321, "Emergency vehicle": 322, "Police car (siren)": 323, "Ambulance (siren)": 324, "Fire engine, fire truck (siren)": 325, "Motorcycle": 326, "Traffic noise, roadway noise": 327, "Rail transport": 328, "Train": 329, "Train whistle": 330, "Train horn": 331, "Railroad car, train wagon": 332, "Train wheels squealing": 333, "Subway, metro, underground": 334, "Aircraft": 335, "Aircraft engine": 336, "Jet engine": 337, "Propeller, airscrew": 338, "Helicopter": 339, "Fixed-wing aircraft, airplane": 340, "Bicycle": 341, "Skateboard": 342, "Engine": 343, "Light engine (high frequency)": 344, "Dental drill, dentist's drill": 345, "Lawn mower": 346, "Chainsaw": 347, "Medium engine (mid frequency)": 348, "Heavy engine (low frequency)": 349, "Engine knocking": 350, "Engine starting": 351, "Idling": 352, "Accelerating, revving, vroom": 353, "Door": 354, "Doorbell": 355, "Ding-dong": 356, "Sliding door": 357, "Slam": 358, "Knock": 359, "Tap": 360, "Squeak": 361, "Cupboard open or close": 362, "Drawer open or close": 363, "Dishes, pots, and pans": 364, "Cutlery, silverware": 365, "Chopping (food)": 366, "Frying (food)": 367, "Microwave oven": 368, "Blender": 369, "Water tap, faucet": 370, "Sink (filling or washing)": 371, "Bathtub (filling or washing)": 372, "Hair dryer": 373, "Toilet flush": 374, "Toothbrush": 375, "Electric toothbrush": 376, "Vacuum cleaner": 377, "Zipper (clothing)": 378, "Keys jangling": 379, "Coin (dropping)": 380, "Scissors": 381, "Electric shaver, electric razor": 382, "Shuffling cards": 383, "Typing": 384, "Typewriter": 385, "Computer keyboard": 386, "Writing": 387, "Alarm": 388, "Telephone": 389, "Telephone bell ringing": 390, "Ringtone": 391, "Telephone dialing, DTMF": 392, "Dial tone": 393, "Busy signal": 394, "Alarm clock": 395, "Siren": 396, "Civil defense siren": 397, "Buzzer": 398, "Smoke detector, smoke alarm": 399, "Fire alarm": 400, "Foghorn": 401, "Whistle": 402, "Steam whistle": 403, "Mechanisms": 404, "Ratchet, pawl": 405, "Clock": 406, "Tick": 407, "Tick-tock": 408, "Gears": 409, "Pulleys": 410, "Sewing machine": 411, "Mechanical fan": 412, "Air conditioning": 413, "Cash register": 414, "Printer": 415, "Camera": 416, "Single-lens reflex camera": 417, "Tools": 418, "Hammer": 419, "Jackhammer": 420, "Sawing": 421, "Filing (rasp)": 422, "Sanding": 423, "Power tool": 424, "Drill": 425, "Explosion": 426, "Gunshot, gunfire": 427, "Machine gun": 428, "Fusillade": 429, "Artillery fire": 430, "Cap gun": 431, "Fireworks": 432, "Firecracker": 433, "Burst, pop": 434, "Eruption": 435, "Boom": 436, "Wood": 437, "Chop": 438, "Splinter": 439, "Crack": 440, "Glass": 441, "Chink, clink": 442, "Shatter": 443, "Liquid": 444, "Splash, splatter": 445, "Slosh": 446, "Squish": 447, "Drip": 448, "Pour": 449, "Trickle, dribble": 450, "Gush": 451, "Fill (with liquid)": 452, "Spray": 453, "Pump (liquid)": 454, "Stir": 455, "Boiling": 456, "Sonar": 457, "Arrow": 458, "Whoosh, swoosh, swish": 459, "Thump, thud": 460, "Thunk": 461, "Electronic tuner": 462, "Effects unit": 463, "Chorus effect": 464, "Basketball bounce": 465, "Bang": 466, "Slap, smack": 467, "Whack, thwack": 468, "Smash, crash": 469, "Breaking": 470, "Bouncing": 471, "Whip": 472, "Flap": 473, "Scratch": 474, "Scrape": 475, "Rub": 476, "Roll": 477, "Crushing": 478, "Crumpling, crinkling": 479, "Tearing": 480, "Beep, bleep": 481, "Ping": 482, "Ding": 483, "Clang": 484, "Squeal": 485, "Creak": 486, "Rustle": 487, "Whir": 488, "Clatter": 489, "Sizzle": 490, "Clicking": 491, "Clickety-clack": 492, "Rumble": 493, "Plop": 494, "Jingle, tinkle": 495, "Hum": 496, "Zing": 497, "Boing": 498, "Crunch": 499, "Silence": 500, "Sine wave": 501, "Harmonic": 502, "Chirp tone": 503, "Sound effect": 504, "Pulse": 505, "Inside, small room": 506, "Inside, large room or hall": 507, "Inside, public space": 508, "Outside, urban or manmade": 509, "Outside, rural or natural": 510, "Reverberation": 511, "Echo": 512, "Noise": 513, "Environmental noise": 514, "Static": 515, "Mains hum": 516, "Distortion": 517, "Sidetone": 518, "Cacophony": 519, "White noise": 520, "Pink noise": 521, "Throbbing": 522, "Vibration": 523, "Television": 524, "Radio": 525, "Field recording": 526}
\ No newline at end of file
+{"Whispering": 0, "Gunshot, gunfire": 1, "Pour": 2, "Wind chime": 3, "Livestock, farm animals, working animals": 4, "Crackle": 5, "Waves, surf": 6, "Chicken, rooster": 7, "Chatter": 8, "Keyboard (musical)": 9, "Bark": 10, "Rail transport": 11, "Gong": 12, "Shatter": 13, "Ratchet, pawl": 14, "Clapping": 15, "Mallet percussion": 16, "Whoosh, swoosh, swish": 17, "Speech synthesizer": 18, "Respiratory sounds": 19, "Sliding door": 20, "Boat, Water vehicle": 21, "Boiling": 22, "Human voice": 23, "Drip": 24, "Thunderstorm": 25, "Male singing": 26, "Sneeze": 27, "Hi-hat": 28, "Guitar": 29, "Crying, sobbing": 30, "Speech": 31, "Slam": 32, "Crack": 33, "Yell": 34, "Drawer open or close": 35, "Run": 36, "Cheering": 37, "Splash, splatter": 38, "Tabla": 39, "Sigh": 40, "Packing tape, duct tape": 41, "Raindrop": 42, "Cymbal": 43, "Fill (with liquid)": 44, "Harp": 45, "Squeak": 46, "Zipper (clothing)": 47, "Tearing": 48, "Alarm": 49, "Skateboard": 50, "Wind instrument, woodwind instrument": 51, "Chink, clink": 52, "Wind": 53, "Ringtone": 54, "Microwave oven": 55, "Power tool": 56, "Dishes, pots, and pans": 57, "Musical instrument": 58, "Door": 59, "Domestic sounds, home sounds": 60, "Subway, metro, underground": 61, "Glockenspiel": 62, "Female speech, woman speaking": 63, "Coin (dropping)": 64, "Mechanical fan": 65, "Male speech, man speaking": 66, "Crowd": 67, "Screech": 68, "Animal": 69, "Human group actions": 70, "Telephone": 71, "Tools": 72, "Giggle": 73, "Crushing": 74, "Thump, thud": 75, "Hammer": 76, "Engine": 77, "Cupboard open or close": 78, "Glass": 79, "Writing": 80, "Clock": 81, "Plucked string instrument": 82, "Fowl": 83, "Water tap, faucet": 84, "Knock": 85, "Trickle, dribble": 86, "Rattle": 87, "Conversation": 88, "Accelerating, revving, vroom": 89, "Fixed-wing aircraft, airplane": 90, "Screaming": 91, "Walk, footsteps": 92, "Stream": 93, "Printer": 94, "Traffic noise, roadway noise": 95, "Motorcycle": 96, "Water": 97, "Scratching (performance technique)": 98, "Tap": 99, "Percussion": 100, "Chuckle, chortle": 101, "Motor vehicle (road)": 102, "Crow": 103, "Vehicle horn, car horn, honking": 104, "Bird vocalization, bird call, bird song": 105, "Drill": 106, "Race car, auto racing": 107, "Meow": 108, "Bass drum": 109, "Drum kit": 110, "Wild animals": 111, "Crash cymbal": 112, "Cough": 113, "Typing": 114, "Bowed string instrument": 115, "Computer keyboard": 116, "Vehicle": 117, "Train": 118, "Applause": 119, "Bicycle": 120, "Tick": 121, "Drum": 122, "Burping, eructation": 123, "Bicycle bell": 124, "Cowbell": 125, "Accordion": 126, "Toilet flush": 127, "Purr": 128, "Church bell": 129, "Cat": 130, "Insect": 131, "Engine starting": 132, "Chewing, mastication": 133, "Sink (filling or washing)": 134, "Dog": 135, "Bird": 136, "Finger snapping": 137, "Child speech, kid speaking": 138, "Wood": 139, "Music": 140, "Sawing": 141, "Bell": 142, "Fireworks": 143, "Crumpling, crinkling": 144, "Ocean": 145, "Gurgling": 146, "Fart": 147, "Mechanisms": 148, "Acoustic guitar": 149, "Singing": 150, "Boom": 151, "Bus": 152, "Cutlery, silverware": 153, "Liquid": 154, "Explosion": 155, "Gull, seagull": 156, "Thunder": 157, "Siren": 158, "Marimba, xylophone": 159, "Female singing": 160, "Tick-tock": 161, "Frog": 162, "Frying (food)": 163, "Buzz": 164, "Car passing by": 165, "Electric guitar": 166, "Gasp": 167, "Rattle (instrument)": 168, "Piano": 169, "Doorbell": 170, "Chime": 171, "Car": 172, "Fire": 173, "Trumpet": 174, "Truck": 175, "Hands": 176, "Domestic animals, pets": 177, "Chirp, tweet": 178, "Breathing": 179, "Cricket": 180, "Tambourine": 181, "Bass guitar": 182, "Idling": 183, "Scissors": 184, "Rain": 185, "Strum": 186, "Shout": 187, "Keys jangling": 188, "Camera": 189, "Hiss": 190, "Growling": 191, "Snare drum": 192, "Brass instrument": 193, "Bathtub (filling or washing)": 194, "Typewriter": 195, "Aircraft": 196, "Organ": 197, "Laughter": 198, "Harmonica": 199}
\ No newline at end of file
diff --git a/class_labels/UrbanSound8K_class_labels_indices.json b/class_labels/UrbanSound8K_class_labels_indices.json
index dc94116..496f37f 100644
--- a/class_labels/UrbanSound8K_class_labels_indices.json
+++ b/class_labels/UrbanSound8K_class_labels_indices.json
@@ -1,1 +1,1 @@
-{"air_conditioner": 0, "car_horn": 1, "children_playing": 2, "dog_bark": 3, "drilling": 4, "engine_idling": 5, "gun_shot": 6, "jackhammer": 7, "siren": 8, "street_music": 9}
\ No newline at end of file
+{"air conditioner": 0, "car horn": 1, "children playing": 2, "dog bark": 3, "drilling": 4, "engine idling": 5, "gun shot": 6, "jackhammer": 7, "siren": 8, "street music": 9}
\ No newline at end of file
diff --git a/class_labels/VGGSound_class_labels_indices.json b/class_labels/VGGSound_class_labels_indices.json
new file mode 100644
index 0000000..d794b2c
--- /dev/null
+++ b/class_labels/VGGSound_class_labels_indices.json
@@ -0,0 +1,1 @@
+{"people crowd": 0, "playing mandolin": 1, "pumping water": 2, "horse neighing": 3, "airplane flyby": 4, "playing drum kit": 5, "pheasant crowing": 6, "duck quacking": 7, "wood thrush calling": 8, "dog bow-wow": 9, "arc welding": 10, "writing on blackboard with chalk": 11, "forging swords": 12, "swimming": 13, "bee, wasp, etc. buzzing": 14, "child singing": 15, "mouse clicking": 16, "playing trombone": 17, "telephone bell ringing": 18, "beat boxing": 19, "cattle mooing": 20, "lions roaring": 21, "ambulance siren": 22, "gibbon howling": 23, "people sniggering": 24, "playing clarinet": 25, "playing bassoon": 26, "playing bongo": 27, "playing electric guitar": 28, "playing badminton": 29, "bull bellowing": 30, "cat caterwauling": 31, "playing sitar": 32, "whale calling": 33, "snake hissing": 34, "people burping": 35, "francolin calling": 36, "fireworks banging": 37, "driving buses": 38, "people belly laughing": 39, "chicken clucking": 40, "playing double bass": 41, "canary calling": 42, "people battle cry": 43, "male singing": 44, "horse clip-clop": 45, "baby crying": 46, "cow lowing": 47, "reversing beeps": 48, "otter growling": 49, "cheetah chirrup": 50, "people running": 51, "ice cream truck, ice cream van": 52, "playing harpsichord": 53, "heart sounds, heartbeat": 54, "pig oinking": 55, "police radio chatter": 56, "cat hissing": 57, "wind chime": 58, "elk bugling": 59, "lions growling": 60, "fly, housefly buzzing": 61, "ferret dooking": 62, "railroad car, train wagon": 63, "church bell ringing": 64, "cat meowing": 65, "wind rustling leaves": 66, "bouncing on trampoline": 67, "mouse squeaking": 68, "sheep bleating": 69, "people eating crisps": 70, "people sneezing": 71, "playing squash": 72, "footsteps on snow": 73, "people humming": 74, "tap dancing": 75, "snake rattling": 76, "elephant trumpeting": 77, "people booing": 78, "disc scratching": 79, "skidding": 80, "cupboard opening or closing": 81, "playing bagpipes": 82, "basketball bounce": 83, "chinchilla barking": 84, "parrot talking": 85, "woodpecker pecking tree": 86, "fire truck siren": 87, "slot machine": 88, "playing french horn": 89, "air conditioning noise": 90, "people finger snapping": 91, "eagle screaming": 92, "playing harmonica": 93, "playing tympani": 94, "zebra braying": 95, "hedge trimmer running": 96, "playing acoustic guitar": 97, "hair dryer drying": 98, "orchestra": 99, "playing darts": 100, "children shouting": 101, "people slurping": 102, "alligators, crocodiles hissing": 103, "mouse pattering": 104, "people marching": 105, "vehicle horn, car horn, honking": 106, "sea lion barking": 107, "people clapping": 108, "hail": 109, "fire crackling": 110, "bathroom ventilation fan running": 111, "opening or closing car doors": 112, "skiing": 113, "dog barking": 114, "race car, auto racing": 115, "subway, metro, underground": 116, "underwater bubbling": 117, "car passing by": 118, "playing tennis": 119, "warbler chirping": 120, "helicopter": 121, "driving motorcycle": 122, "train wheels squealing": 123, "baby laughter": 124, "driving snowmobile": 125, "bird squawking": 126, "cuckoo bird calling": 127, "people whistling": 128, "shot football": 129, "playing tuning fork": 130, "dog howling": 131, "playing violin, fiddle": 132, "people eating": 133, "baltimore oriole calling": 134, "playing timbales": 135, "door slamming": 136, "people shuffling": 137, "typing on typewriter": 138, "magpie calling": 139, "playing harp": 140, "playing hammond organ": 141, "people eating apple": 142, "mosquito buzzing": 143, "playing oboe": 144, "playing volleyball": 145, "using sewing machines": 146, "electric grinder grinding": 147, "cutting hair with electric trimmers": 148, "splashing water": 149, "people sobbing": 150, "female singing": 151, "wind noise": 152, "car engine knocking": 153, "black capped chickadee calling": 154, "people screaming": 155, "cat growling": 156, "penguins braying": 157, "people coughing": 158, "metronome": 159, "train horning": 160, "goat bleating": 161, "playing tambourine": 162, "fox barking": 163, "airplane": 164, "firing cannon": 165, "thunder": 166, "smoke detector beeping": 167, "playing erhu": 168, "ice cracking": 169, "dog growling": 170, "playing saxophone": 171, "owl hooting": 172, "playing trumpet": 173, "sailing": 174, "waterfall burbling": 175, "machine gun shooting": 176, "baby babbling": 177, "playing synthesizer": 178, "donkey, ass braying": 179, "people cheering": 180, "playing shofar": 181, "playing hockey": 182, "playing banjo": 183, "cricket chirping": 184, "playing snare drum": 185, "ripping paper": 186, "child speech, kid speaking": 187, "crow cawing": 188, "sloshing water": 189, "playing zither": 190, "scuba diving": 191, "playing steelpan": 192, "goose honking": 193, "tapping guitar": 194, "spraying water": 195, "playing bass drum": 196, "printer printing": 197, "playing ukulele": 198, "ocean burbling": 199, "playing didgeridoo": 200, "sharpen knife": 201, "typing on computer keyboard": 202, "playing table tennis": 203, "rope skipping": 204, "playing marimba, xylophone": 205, "playing bugle": 206, "playing guiro": 207, "playing flute": 208, "tornado roaring": 209, "stream burbling": 210, "electric shaver, electric razor shaving": 211, "playing gong": 212, "eating with cutlery": 213, "playing piano": 214, "people giggling": 215, "chicken crowing": 216, "female speech, woman speaking": 217, "golf driving": 218, "frog croaking": 219, "people eating noodle": 220, "mynah bird singing": 221, "playing timpani": 222, "playing congas": 223, "dinosaurs bellowing": 224, "playing bass guitar": 225, "turkey gobbling": 226, "chipmunk chirping": 227, "chopping food": 228, "striking bowling": 229, "missile launch": 230, "squishing water": 231, "civil defense siren": 232, "blowtorch igniting": 233, "tractor digging": 234, "lighting firecrackers": 235, "playing theremin": 236, "train whistling": 237, "people nose blowing": 238, "car engine starting": 239, "lathe spinning": 240, "playing cello": 241, "motorboat, speedboat acceleration": 242, "playing vibraphone": 243, "playing washboard": 244, "playing cornet": 245, "pigeon, dove cooing": 246, "roller coaster running": 247, "opening or closing car electric windows": 248, "foghorn": 249, "coyote howling": 250, "hammering nails": 251, "toilet flushing": 252, "strike lighter": 253, "bird wings flapping": 254, "playing steel guitar, slide guitar": 255, "volcano explosion": 256, "people whispering": 257, "bowling impact": 258, "yodelling": 259, "firing muskets": 260, "raining": 261, "singing bowl": 262, "plastic bottle crushing": 263, "chimpanzee pant-hooting": 264, "playing electronic organ": 265, "chainsawing trees": 266, "dog baying": 267, "lawn mowing": 268, "people babbling": 269, "striking pool": 270, "eletric blender running": 271, "playing tabla": 272, "cap gun shooting": 273, "planing timber": 274, "air horn": 275, "sliding door": 276, "cell phone buzzing": 277, "sea waves": 278, "playing castanets": 279, "singing choir": 280, "people slapping": 281, "barn swallow calling": 282, "people hiccup": 283, "vacuum cleaner cleaning floors": 284, "playing lacrosse": 285, "bird chirping, tweeting": 286, "lip smacking": 287, "chopping wood": 288, "police car (siren)": 289, "running electric fan": 290, "cattle, bovinae cowbell": 291, "people gargling": 292, "opening or closing drawers": 293, "playing djembe": 294, "skateboarding": 295, "cat purring": 296, "rowboat, canoe, kayak rowing": 297, "engine accelerating, revving, vroom": 298, "playing glockenspiel": 299, "popping popcorn": 300, "car engine idling": 301, "alarm clock ringing": 302, "dog whimpering": 303, "playing accordion": 304, "playing cymbal": 305, "male speech, man speaking": 306, "rapping": 307, "people farting": 308}
\ No newline at end of file
diff --git a/class_labels/audioset_fsd50k_class_labels_indices.json b/class_labels/audioset_fsd50k_class_labels_indices.json
new file mode 100644
index 0000000..fd8014e
--- /dev/null
+++ b/class_labels/audioset_fsd50k_class_labels_indices.json
@@ -0,0 +1,1 @@
+{"Speech": 0, "Male speech, man speaking": 1, "Female speech, woman speaking": 2, "Child speech, kid speaking": 3, "Conversation": 4, "Narration, monologue": 5, "Babbling": 6, "Speech synthesizer": 7, "Shout": 8, "Bellow": 9, "Whoop": 10, "Yell": 11, "Battle cry": 12, "Children shouting": 13, "Screaming": 14, "Whispering": 15, "Laughter": 16, "Baby laughter": 17, "Giggle": 18, "Snicker": 19, "Belly laugh": 20, "Chuckle, chortle": 21, "Crying, sobbing": 22, "Baby cry, infant cry": 23, "Whimper": 24, "Wail, moan": 25, "Sigh": 26, "Singing": 27, "Choir": 28, "Yodeling": 29, "Chant": 30, "Mantra": 31, "Male singing": 32, "Female singing": 33, "Child singing": 34, "Synthetic singing": 35, "Rapping": 36, "Humming": 37, "Groan": 38, "Grunt": 39, "Whistling": 40, "Breathing": 41, "Wheeze": 42, "Snoring": 43, "Gasp": 44, "Pant": 45, "Snort": 46, "Cough": 47, "Throat clearing": 48, "Sneeze": 49, "Sniff": 50, "Run": 51, "Shuffle": 52, "Walk, footsteps": 53, "Chewing, mastication": 54, "Biting": 55, "Gargling": 56, "Stomach rumble": 57, "Burping, eructation": 58, "Hiccup": 59, "Fart": 60, "Hands": 61, "Finger snapping": 62, "Clapping": 63, "Heart sounds, heartbeat": 64, "Heart murmur": 65, "Cheering": 66, "Applause": 67, "Chatter": 68, "Crowd": 69, "Hubbub, speech noise, speech babble": 70, "Children playing": 71, "Animal": 72, "Domestic animals, pets": 73, "Dog": 74, "Bark": 75, "Yip": 76, "Howl": 77, "Bow-wow": 78, "Growling": 79, "Whimper (dog)": 80, "Cat": 81, "Purr": 82, "Meow": 83, "Hiss": 84, "Caterwaul": 85, "Livestock, farm animals, working animals": 86, "Horse": 87, "Clip-clop": 88, "Neigh, whinny": 89, "Cattle, bovinae": 90, "Moo": 91, "Cowbell": 92, "Pig": 93, "Oink": 94, "Goat": 95, "Bleat": 96, "Sheep": 97, "Fowl": 98, "Chicken, rooster": 99, "Cluck": 100, "Crowing, cock-a-doodle-doo": 101, "Turkey": 102, "Gobble": 103, "Duck": 104, "Quack": 105, "Goose": 106, "Honk": 107, "Wild animals": 108, "Roaring cats (lions, tigers)": 109, "Roar": 110, "Bird": 111, "Bird vocalization, bird call, bird song": 112, "Chirp, tweet": 113, "Squawk": 114, "Pigeon, dove": 115, "Coo": 116, "Crow": 117, "Caw": 118, "Owl": 119, "Hoot": 120, "Bird flight, flapping wings": 121, "Canidae, dogs, wolves": 122, "Rodents, rats, mice": 123, "Mouse": 124, "Patter": 125, "Insect": 126, "Cricket": 127, "Mosquito": 128, "Fly, housefly": 129, "Buzz": 130, "Bee, wasp, etc.": 131, "Frog": 132, "Croak": 133, "Snake": 134, "Rattle": 135, "Whale vocalization": 136, "Music": 137, "Musical instrument": 138, "Plucked string instrument": 139, "Guitar": 140, "Electric guitar": 141, "Bass guitar": 142, "Acoustic guitar": 143, "Steel guitar, slide guitar": 144, "Tapping (guitar technique)": 145, "Strum": 146, "Banjo": 147, "Sitar": 148, "Mandolin": 149, "Zither": 150, "Ukulele": 151, "Keyboard (musical)": 152, "Piano": 153, "Electric piano": 154, "Organ": 155, "Electronic organ": 156, "Hammond organ": 157, "Synthesizer": 158, "Sampler": 159, "Harpsichord": 160, "Percussion": 161, "Drum kit": 162, "Drum machine": 163, "Drum": 164, "Snare drum": 165, "Rimshot": 166, "Drum roll": 167, "Bass drum": 168, "Timpani": 169, "Tabla": 170, "Cymbal": 171, "Hi-hat": 172, "Wood block": 173, "Tambourine": 174, "Rattle (instrument)": 175, "Maraca": 176, "Gong": 177, "Tubular bells": 178, "Mallet percussion": 179, "Marimba, xylophone": 180, "Glockenspiel": 181, "Vibraphone": 182, "Steelpan": 183, "Orchestra": 184, "Brass instrument": 185, "French horn": 186, "Trumpet": 187, "Trombone": 188, "Bowed string instrument": 189, "String section": 190, "Violin, fiddle": 191, "Pizzicato": 192, "Cello": 193, "Double bass": 194, "Wind instrument, woodwind instrument": 195, "Flute": 196, "Saxophone": 197, "Clarinet": 198, "Harp": 199, "Bell": 200, "Church bell": 201, "Jingle bell": 202, "Bicycle bell": 203, "Tuning fork": 204, "Chime": 205, "Wind chime": 206, "Change ringing (campanology)": 207, "Harmonica": 208, "Accordion": 209, "Bagpipes": 210, "Didgeridoo": 211, "Shofar": 212, "Theremin": 213, "Singing bowl": 214, "Scratching (performance technique)": 215, "Pop music": 216, "Hip hop music": 217, "Beatboxing": 218, "Rock music": 219, "Heavy metal": 220, "Punk rock": 221, "Grunge": 222, "Progressive rock": 223, "Rock and roll": 224, "Psychedelic rock": 225, "Rhythm and blues": 226, "Soul music": 227, "Reggae": 228, "Country": 229, "Swing music": 230, "Bluegrass": 231, "Funk": 232, "Folk music": 233, "Middle Eastern music": 234, "Jazz": 235, "Disco": 236, "Classical music": 237, "Opera": 238, "Electronic music": 239, "House music": 240, "Techno": 241, "Dubstep": 242, "Drum and bass": 243, "Electronica": 244, "Electronic dance music": 245, "Ambient music": 246, "Trance music": 247, "Music of Latin America": 248, "Salsa music": 249, "Flamenco": 250, "Blues": 251, "Music for children": 252, "New-age music": 253, "Vocal music": 254, "A capella": 255, "Music of Africa": 256, "Afrobeat": 257, "Christian music": 258, "Gospel music": 259, "Music of Asia": 260, "Carnatic music": 261, "Music of Bollywood": 262, "Ska": 263, "Traditional music": 264, "Independent music": 265, "Song": 266, "Background music": 267, "Theme music": 268, "Jingle (music)": 269, "Soundtrack music": 270, "Lullaby": 271, "Video game music": 272, "Christmas music": 273, "Dance music": 274, "Wedding music": 275, "Happy music": 276, "Funny music": 277, "Sad music": 278, "Tender music": 279, "Exciting music": 280, "Angry music": 281, "Scary music": 282, "Wind": 283, "Rustling leaves": 284, "Wind noise (microphone)": 285, "Thunderstorm": 286, "Thunder": 287, "Water": 288, "Rain": 289, "Raindrop": 290, "Rain on surface": 291, "Stream": 292, "Waterfall": 293, "Ocean": 294, "Waves, surf": 295, "Steam": 296, "Gurgling": 297, "Fire": 298, "Crackle": 299, "Vehicle": 300, "Boat, Water vehicle": 301, "Sailboat, sailing ship": 302, "Rowboat, canoe, kayak": 303, "Motorboat, speedboat": 304, "Ship": 305, "Motor vehicle (road)": 306, "Car": 307, "Vehicle horn, car horn, honking": 308, "Toot": 309, "Car alarm": 310, "Power windows, electric windows": 311, "Skidding": 312, "Tire squeal": 313, "Car passing by": 314, "Race car, auto racing": 315, "Truck": 316, "Air brake": 317, "Air horn, truck horn": 318, "Reversing beeps": 319, "Ice cream truck, ice cream van": 320, "Bus": 321, "Emergency vehicle": 322, "Police car (siren)": 323, "Ambulance (siren)": 324, "Fire engine, fire truck (siren)": 325, "Motorcycle": 326, "Traffic noise, roadway noise": 327, "Rail transport": 328, "Train": 329, "Train whistle": 330, "Train horn": 331, "Railroad car, train wagon": 332, "Train wheels squealing": 333, "Subway, metro, underground": 334, "Aircraft": 335, "Aircraft engine": 336, "Jet engine": 337, "Propeller, airscrew": 338, "Helicopter": 339, "Fixed-wing aircraft, airplane": 340, "Bicycle": 341, "Skateboard": 342, "Engine": 343, "Light engine (high frequency)": 344, "Dental drill, dentist's drill": 345, "Lawn mower": 346, "Chainsaw": 347, "Medium engine (mid frequency)": 348, "Heavy engine (low frequency)": 349, "Engine knocking": 350, "Engine starting": 351, "Idling": 352, "Accelerating, revving, vroom": 353, "Door": 354, "Doorbell": 355, "Ding-dong": 356, "Sliding door": 357, "Slam": 358, "Knock": 359, "Tap": 360, "Squeak": 361, "Cupboard open or close": 362, "Drawer open or close": 363, "Dishes, pots, and pans": 364, "Cutlery, silverware": 365, "Chopping (food)": 366, "Frying (food)": 367, "Microwave oven": 368, "Blender": 369, "Water tap, faucet": 370, "Sink (filling or washing)": 371, "Bathtub (filling or washing)": 372, "Hair dryer": 373, "Toilet flush": 374, "Toothbrush": 375, "Electric toothbrush": 376, "Vacuum cleaner": 377, "Zipper (clothing)": 378, "Keys jangling": 379, "Coin (dropping)": 380, "Scissors": 381, "Electric shaver, electric razor": 382, "Shuffling cards": 383, "Typing": 384, "Typewriter": 385, "Computer keyboard": 386, "Writing": 387, "Alarm": 388, "Telephone": 389, "Telephone bell ringing": 390, "Ringtone": 391, "Telephone dialing, DTMF": 392, "Dial tone": 393, "Busy signal": 394, "Alarm clock": 395, "Siren": 396, "Civil defense siren": 397, "Buzzer": 398, "Smoke detector, smoke alarm": 399, "Fire alarm": 400, "Foghorn": 401, "Whistle": 402, "Steam whistle": 403, "Mechanisms": 404, "Ratchet, pawl": 405, "Clock": 406, "Tick": 407, "Tick-tock": 408, "Gears": 409, "Pulleys": 410, "Sewing machine": 411, "Mechanical fan": 412, "Air conditioning": 413, "Cash register": 414, "Printer": 415, "Camera": 416, "Single-lens reflex camera": 417, "Tools": 418, "Hammer": 419, "Jackhammer": 420, "Sawing": 421, "Filing (rasp)": 422, "Sanding": 423, "Power tool": 424, "Drill": 425, "Explosion": 426, "Gunshot, gunfire": 427, "Machine gun": 428, "Fusillade": 429, "Artillery fire": 430, "Cap gun": 431, "Fireworks": 432, "Firecracker": 433, "Burst, pop": 434, "Eruption": 435, "Boom": 436, "Wood": 437, "Chop": 438, "Splinter": 439, "Crack": 440, "Glass": 441, "Chink, clink": 442, "Shatter": 443, "Liquid": 444, "Splash, splatter": 445, "Slosh": 446, "Squish": 447, "Drip": 448, "Pour": 449, "Trickle, dribble": 450, "Gush": 451, "Fill (with liquid)": 452, "Spray": 453, "Pump (liquid)": 454, "Stir": 455, "Boiling": 456, "Sonar": 457, "Arrow": 458, "Whoosh, swoosh, swish": 459, "Thump, thud": 460, "Thunk": 461, "Electronic tuner": 462, "Effects unit": 463, "Chorus effect": 464, "Basketball bounce": 465, "Bang": 466, "Slap, smack": 467, "Whack, thwack": 468, "Smash, crash": 469, "Breaking": 470, "Bouncing": 471, "Whip": 472, "Flap": 473, "Scratch": 474, "Scrape": 475, "Rub": 476, "Roll": 477, "Crushing": 478, "Crumpling, crinkling": 479, "Tearing": 480, "Beep, bleep": 481, "Ping": 482, "Ding": 483, "Clang": 484, "Squeal": 485, "Creak": 486, "Rustle": 487, "Whir": 488, "Clatter": 489, "Sizzle": 490, "Clicking": 491, "Clickety-clack": 492, "Rumble": 493, "Plop": 494, "Jingle, tinkle": 495, "Hum": 496, "Zing": 497, "Boing": 498, "Crunch": 499, "Silence": 500, "Sine wave": 501, "Harmonic": 502, "Chirp tone": 503, "Sound effect": 504, "Pulse": 505, "Inside, small room": 506, "Inside, large room or hall": 507, "Inside, public space": 508, "Outside, urban or manmade": 509, "Outside, rural or natural": 510, "Reverberation": 511, "Echo": 512, "Noise": 513, "Environmental noise": 514, "Static": 515, "Mains hum": 516, "Distortion": 517, "Sidetone": 518, "Cacophony": 519, "White noise": 520, "Pink noise": 521, "Throbbing": 522, "Vibration": 523, "Television": 524, "Radio": 525, "Field recording": 526, "Respiratory sounds": 527, "Human voice": 528, "Packing tape, duct tape": 529, "Domestic sounds, home sounds": 530, "Screech": 531, "Human group actions": 532, "Crash cymbal": 533, "Gull, seagull": 534}
\ No newline at end of file
diff --git a/run_code.txt b/run_code.txt
index 3b79f8b..03b9632 100644
--- a/run_code.txt
+++ b/run_code.txt
@@ -20,7 +20,7 @@ CUDA_VISIBLE_DEVICES=0 python -m training.main \
     --datasetinfos "train" "unbalanced_train" "balanced_train"
 
 # Toy dataset Linear Probe
-CUDA_VISIBLE_DEVICES=0 python -m training.lp_main \
+CUDA_VISIBLE_DEVICES=0 python -m training.main \
     --save-frequency 1 \
     --dataset-type="toy" \
     --train-data="/home/la/kechen/Research/ke_zsasp/workspace/hdf5s/indexes/balanced_train.h5"  \
@@ -29,11 +29,11 @@ CUDA_VISIBLE_DEVICES=0 python -m training.lp_main \
     --val-ipc="/home/la/kechen/Research/clap/CLAP/idc_eval.npy" \
     --precision="fp32" \
     --warmup 10000 \
-    --batch-size=32 \
+    --batch-size=4 \
     --lr=1e-3 \
     --wd=0.1 \
     --epochs=1 \
-    --workers=4 \
+    --workers=0 \
     --amodel HTSAT-tiny \
     --tmodel transformer \
     --no-eval \
@@ -400,4 +400,34 @@ conda activate ~/audio_clip
 cd ~/CLAP/src
 
 num_nodes=16
-jsrun --launch_distribution=packed -bpacked:7 -n${num_nodes} -r1 -a6 -g6 -c42 ./launch-smt4.sh ./run_summit_clap.sh
\ No newline at end of file
+jsrun --launch_distribution=packed -bpacked:7 -n${num_nodes} -r1 -a6 -g6 -c42 ./launch-smt4.sh ./run_summit_clap.sh
+
+num_nodes=2
+jsrun --launch_distribution=packed -bpacked:7 -n${num_nodes} -r1 -a6 -g6 -c42 ./launch-smt4.sh ./run_summit_clap.sh
+
+jsrun --launch_distribution=packed -bpacked:7 -n1 -r1 -a1 -g1 -c7 ./launch-smt1.sh ./run_summit_clap_eval.sh
+
+num_nodes=1
+jsrun --launch_distribution=packed -bpacked:7 -n${num_nodes} -r1 -a6 -g6 -c42 ./launch-smt4.sh ./run_summit_clap.sh
+
+num_nodes=2
+jsrun --launch_distribution=packed -bpacked:7 -n${num_nodes} -r1 -a6 -g6 -c42 ./launch-smt4.sh ./run_summit_clap.sh
+
+jsrun --launch_distribution=packed -bpacked:7 -n1 -r1 -a1 -g1 -c7 ./launch-smt1.sh ./run_summit_clap_eval.sh
+
+num_nodes=1
+jsrun --launch_distribution=packed -bpacked:7 -n${num_nodes} -r1 -a6 -g6 -c42 ./launch-smt4.sh ./run_summit_clap.sh
+
+
+# pretrained_audio = ''
+# audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+# audio_ckpt = audio_ckpt['state_dict']
+# save_ckpt = {}
+# keys = list(audio_ckpt.keys())
+# for key in keys:
+#     if 'audio_branch' in key:
+#         v = audio_ckpt.pop(key)
+#         save_key = key.replace('module.clap_model.', '')
+#         save_ckpt[save_key] = v
+#
+# torch.save(save_ckpt, '10.13_pann_epoch_50_mixup_processed.pt')
\ No newline at end of file
diff --git a/src/check_ckpt.py b/src/check_ckpt.py
index 1fd0c23..d034dd4 100644
--- a/src/check_ckpt.py
+++ b/src/check_ckpt.py
@@ -40,6 +40,10 @@ def check_ckpt_diff(ckpt_a, ckpt_b, key_include=None, key_exclude=None, device='
 # Transformer no freeze:
 # check_ckpt_diff("/fsx/clap_logs/2022_09_11-19_37_08-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_10.pt", "/fsx/clap_logs/2022_09_11-19_37_08-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_100.pt", "text_branch.resblocks")
 
+check_ckpt_diff("/fsx/clap_logs/2022_09_29-23_42_40-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_1.pt",
+                "/fsx/clap_logs/2022_09_29-23_42_40-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_2.pt",
+                "text_branch.resblocks")
+
 # key module.text_branch.resblocks.0.attn.in_proj_weight is different
 # key module.text_branch.resblocks.0.attn.in_proj_bias is different
 # key module.text_branch.resblocks.0.attn.out_proj.weight is different
diff --git a/src/evaluate/eval_dcase.py b/src/evaluate/eval_dcase.py
index a32ee02..910bfe7 100644
--- a/src/evaluate/eval_dcase.py
+++ b/src/evaluate/eval_dcase.py
@@ -10,17 +10,20 @@
 import librosa
 from tqdm import tqdm
 import numpy as np
+import os
+from training.params import parse_args
 
 
-def get_output_from_single_audio(audio, text, model):
-    audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding")['embedding']
-    if audio_embedding.ndim > 1:
-        audio_embedding = audio_embedding.mean(dim=0, keepdim=True)
-    else:
-        audio_embedding = audio_embedding.unsqueeze(0)
-    audio_features = model.audio_projection(audio_embedding)
+def get_output_from_single_audio(audio, text, model, device):
+
+    # audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding", device=device)['embedding']
+    # if audio_embedding.ndim > 1:
+    #     audio_embedding = audio_embedding.mean(dim=0, keepdim=True)
+    # else:
+    #     audio_embedding = audio_embedding.unsqueeze(0)
+    audio_features = model(audio, None, device)
     audio_features = F.normalize(audio_features, dim=-1)
-    text_features = model.encode_text(text)
+    text_features = model(None, text, device=device)
     text_features = F.normalize(text_features, dim=-1)
 
     # CHANGE: before normalize or after
@@ -48,159 +51,14 @@ def get_metrics(text_to_audio_logits):
 
 
 if __name__ == '__main__':
+    args = parse_args()
 
-    # top-2 transformer + top-1 CNN
-    model_path_all = [
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_2.pt',
-    ]
-
-    model_type_all = [
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-    ]
-
-    # all transformer + top-1 CNN
-    """
-    model_path_all = [
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_2.pt',
-    ]
-
-    model_type_all = [
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-    ]
-    """
-
-    # all transformers
-    """
-    model_path_all = [
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-    ]
+    model_path = args.pretrained
 
-    model_type_all = [
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-    ]
-    """
-
-    # all systems
-    """
-    model_path_all = [
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-21_01_57-model_PANN-14-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-21_01_57-model_PANN-14-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-21_01_57-model_PANN-14-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-    ]
-
-    model_type_all = [
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-    ]
-    """
-
-
-    clotho_test_preprocessed_dir = "/mnt/audio_clip/dcase_submission_systems/clotho_test/"
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--device', type=str, default='cpu')
-    args = parser.parse_args()
+    clotho_test_preprocessed_dir = "/fsx/yusong/clotho_test_set/test"
 
     cudnn.benchmark = True
     cudnn.deterministic = False
-    pretrained = 'openai'
 
     audio_features_ensemble_all = []
     text_features_ensemble_all = []
@@ -209,95 +67,86 @@ def get_metrics(text_to_audio_logits):
     logit_scale_a_ensemble_all = []
     logit_scale_t_ensemble_all = []
 
-    for model_path, model_type in zip(model_path_all, model_type_all):
-        device = torch.device(args.device)
-        model, model_cfg = create_model(
-            model_type,
-            pretrained,
-            precision='fp32',
-            device=device,
-            jit=False,
-            force_quick_gelu=False,
-        )
-
-        # load model
-        checkpoint = torch.load(model_path, map_location=device)
-        if "epoch" in checkpoint:
-            # resuming a train checkpoint w/ epoch and optimizer state
-            start_epoch = checkpoint["epoch"]
-            sd = checkpoint["state_dict"]
-            if next(iter(sd.items()))[0].startswith(
-                    "module"
-            ):
-                sd = {k[len("module."):]: v for k, v in sd.items()}
-            model.load_state_dict(sd)
-        else:
-            # loading a bare (model only) checkpoint for fine-tune or evaluation
-            model.load_state_dict(checkpoint)
 
-        model.to(device)
-        model.eval()
-        for param in model.parameters():
-            param.requires_grad = False
-
-        # take every 5th file because clotho has 5 texts for 1 audio
-        test_file_list = sorted(glob.glob(f"{clotho_test_preprocessed_dir}/*.flac"))[::5]
-
-        audio_features_all = []
-        text_features_all = []
-        audio_features_mlp_all = []
-        text_features_mlp_all = []
-        logit_scale_a_all = []
-        logit_scale_t_all = []
-
-        with torch.no_grad():
-            for file_path in tqdm(test_file_list):
-                json_path = file_path.replace(".flac", ".json")
-                with open(json_path, "r") as f:
-                    json_data = json.load(f)
-                audio, sr = librosa.load(file_path, sr=48000, mono=True)
-                audio = torch.from_numpy(audio).to(device)
-                text = json_data["original_data"]["all_captions"]
-                text = tokenize(text).to(device)
-
-                audio_features, text_features, audio_features_mlp, text_features_mlp, logit_scale_a, logit_scale_t = \
-                    get_output_from_single_audio(audio, text, model)
-
-                audio_features_all.append(audio_features.detach().cpu())
-                text_features_all.append(text_features.detach().cpu())
-                audio_features_mlp_all.append(audio_features_mlp.detach().cpu())
-                text_features_mlp_all.append(text_features_mlp.detach().cpu())
-                logit_scale_a_all.append(logit_scale_a.detach().cpu())
-                logit_scale_t_all.append(logit_scale_t.detach().cpu())
-
-        audio_features_ensemble_all.append(torch.cat(audio_features_all))
-        text_features_ensemble_all.append(torch.cat(text_features_all))
-        audio_features_mlp_ensemble_all.append(torch.cat(audio_features_mlp_all))
-        text_features_mlp_ensemble_all.append(torch.cat(text_features_mlp_all))
-        logit_scale_a_ensemble_all.append(logit_scale_a_all[0])
-        logit_scale_t_ensemble_all.append(logit_scale_t_all[0])
-
-    text_to_audio_logits_ensemble_all = []
-    for i in range(len(model_path_all)):
-        audio_features = audio_features_ensemble_all[i]
-        text_features = text_features_ensemble_all[i]
-        audio_features_mlp = audio_features_mlp_ensemble_all[i]
-        text_features_mlp = text_features_mlp_ensemble_all[i]
-        logit_scale_a = logit_scale_a_ensemble_all[i]
-        logit_scale_t = logit_scale_t_ensemble_all[i]
-
-        a_logits_per_audio = (logit_scale_a * audio_features @ text_features_mlp.t()).detach().cpu()
-        a_logits_per_text = a_logits_per_audio.t().detach().cpu()
-        t_logits_per_audio = (logit_scale_t * audio_features_mlp @ text_features.t()).detach().cpu()
-        t_logits_per_text = t_logits_per_audio.t().detach().cpu()
+    device = torch.device('cuda')
+    model, clap_model_cfg = create_model(
+        args.amodel,
+        args.tmodel,
+        args.pretrained,
+        precision=args.precision,
+        device=device,
+        jit=args.torchscript,
+        force_quick_gelu=args.force_quick_gelu,
+        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+        skip_params=False
+    )
 
-        text_to_audio_logits = (a_logits_per_text + t_logits_per_text) / 2
-        text_to_audio_logits_ensemble_all.append(text_to_audio_logits)
+    # load model
+    checkpoint = torch.load(model_path, map_location=device)
+    if "epoch" in checkpoint:
+        # resuming a train checkpoint w/ epoch and optimizer state
+        start_epoch = checkpoint["epoch"]
+        sd = checkpoint["state_dict"]
+        if next(iter(sd.items()))[0].startswith(
+                "module"
+        ):
+            sd = {k[len("module."):]: v for k, v in sd.items()}
+        model.load_state_dict(sd)
+    else:
+        # loading a bare (model only) checkpoint for fine-tune or evaluation
+        model.load_state_dict(checkpoint)
+
+    model.to(device)
+    model.eval()
+    for param in model.parameters():
+        param.requires_grad = False
+
+    # take every 5th file because clotho has 5 texts for 1 audio
+    test_file_list = sorted(glob.glob(f"{clotho_test_preprocessed_dir}/*.flac"))
+
+    audio_features_all = []
+    text_features_all = []
+    audio_features_mlp_all = []
+    text_features_mlp_all = []
+    logit_scale_a_all = []
+    logit_scale_t_all = []
+
+    with torch.no_grad():
+        for file_path in tqdm(test_file_list):
+            json_path = file_path.replace(".flac", ".json")
+            with open(json_path, "r") as f:
+                json_data = json.load(f)
+            audio, sr = librosa.load(file_path, sr=48000, mono=True)
+            audio = torch.from_numpy(audio).to(device)
+            audio = {'waveform': audio.unsqueeze(0), 'sample_rate': sr}
+            text = json_data["text"]
+
+            if args.tmodel == "transformer":
+                from open_clip import tokenize
+                text = tokenize(text)
+            else:
+                from training.data import tokenizer
+                text = tokenizer(text)  # 5 texts for each audio
+
+            audio_features, text_features, audio_features_mlp, text_features_mlp, logit_scale_a, logit_scale_t = \
+                get_output_from_single_audio(audio, text, model, device)
+
+            audio_features_all.append(audio_features.detach().cpu())
+            text_features_all.append(text_features.detach().cpu())
+            audio_features_mlp_all.append(audio_features_mlp.detach().cpu())
+            text_features_mlp_all.append(text_features_mlp.detach().cpu())
+            logit_scale_a_all.append(logit_scale_a.detach().cpu())
+            logit_scale_t_all.append(logit_scale_t.detach().cpu())
+
+    audio_features = torch.cat(audio_features_all)
+    text_features = torch.cat(text_features_all)
+    logit_scale_a = logit_scale_a_all[0]
+
+    logits_per_audio = (logit_scale_a * audio_features @ text_features.t()).detach().cpu()
+    logits_per_text = logits_per_audio.t().detach().cpu()
 
     metrics = get_metrics(
-        torch.stack(text_to_audio_logits_ensemble_all).mean(dim=0)
+        logits_per_text
     )
 
     print(metrics)
-
-    # TODO:[Tianyu]  load from DCASE contest dataset, get the retrieval file, and save csv file
diff --git a/src/evaluate/eval_linear_probe.py b/src/evaluate/eval_linear_probe.py
index cc49dc7..bf5a524 100644
--- a/src/evaluate/eval_linear_probe.py
+++ b/src/evaluate/eval_linear_probe.py
@@ -177,11 +177,19 @@ def main():
         os.makedirs(log_base_path, exist_ok=True)
         log_filename = f"out-{args.rank}" if args.log_local else "out.log"
         args.log_path = os.path.join(log_base_path, log_filename)
-        if os.path.exists(args.log_path):
-            print(
-                "Error. Experiment already exists. Use --name {} to specify a new experiment."
-            )
-            return -1
+
+        # avoid log dir in same name:
+        postfix = 0
+        while os.path.exists(args.log_path):
+            postfix += 1
+            log_base_path_new = log_base_path+'-'+str(postfix)
+            os.makedirs(log_base_path_new, exist_ok=True)
+            log_filename = f"out-{args.rank}" if args.log_local else "out.log"
+            args.log_path = os.path.join(log_base_path_new, log_filename)
+            # print(
+            #     "Error. Experiment already exists. Use --name {} to specify a new experiment."
+            # )
+            # return -1
 
     # Set logger
     args.log_level = logging.DEBUG if args.debug else logging.INFO
@@ -297,7 +305,9 @@ def lp_main(args, device, writer, pretrain_epoch, idx):
         jit=args.torchscript,
         force_quick_gelu=args.force_quick_gelu,
         openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-        skip_params=False
+        skip_params=False,
+        enable_fusion=args.enable_fusion,
+        fusion_type=args.fusion_type
     )
     
     args.lp_out_ch = len(list(args.class_index_dict.keys()))
diff --git a/src/evaluate/eval_retrieval_main.py b/src/evaluate/eval_retrieval_main.py
index 8356b26..002f2eb 100644
--- a/src/evaluate/eval_retrieval_main.py
+++ b/src/evaluate/eval_retrieval_main.py
@@ -31,18 +31,85 @@ def find_params_value(file, key):
     return None
 
 
+def evaluate_zeroshot(model, data, start_epoch, args, writer):
+    dataloader = data["val"].dataloader
+    metrics = {}
+    device = torch.device(args.device)
+    model.eval()
+    metrics.update({"epoch": start_epoch})
+
+    all_audio_features = []
+    all_class_labels = []
+    with torch.no_grad():
+        for i, batch in enumerate(dataloader):
+            audios = batch  # contains mel_spec, wavform, and longer list
+            audio_features = model(audios, None, device)
+            audio_features = F.normalize(audio_features, dim=-1)
+            all_audio_features.append(audio_features.detach().cpu())
+            all_class_labels.append(torch.argmax(batch["class_label"], 1).long())
+        all_audio_features = torch.cat(all_audio_features, dim=0)
+        all_class_labels = torch.cat(all_class_labels, dim=0)
+        metrics["num_samples"] = all_audio_features.shape[0]
+
+        # get text features
+        all_texts = ["This is a sound of " + t for t in args.class_index_dict.keys()]
+        # (yusong): a hack, can make it better
+        if args.tmodel == "transformer":
+            from open_clip import tokenize
+            all_texts = tokenize(all_texts)
+        else:
+            from training.data import tokenizer
+            all_texts = tokenizer(all_texts)
+        all_text_features = model(None, all_texts, device)
+        all_text_features = F.normalize(all_text_features, dim=-1).detach().cpu()
+
+        # compute similarity
+        logit_scale_a, logit_scale_t = model(None, None, device)
+        logit_scale_a = logit_scale_a.cpu()
+
+        logits_per_audio = (logit_scale_a * all_audio_features @ all_text_features.t()).detach().cpu()
+        logits_per_text = logits_per_audio.t().detach().cpu()
+
+        ground_truth = all_class_labels.view(-1, 1)
+        logit = logits_per_audio
+
+        ranking = torch.argsort(logit, descending=True)
+        preds = torch.where(ranking == ground_truth)[1]  # (yusong) this line is slow because it uses single thread
+        preds = preds.detach().cpu().numpy()
+        metrics[f"{args.datasetnames[0]}_mean_rank"] = preds.mean() + 1
+        metrics[f"{args.datasetnames[0]}_median_rank"] = np.floor(np.median(preds)) + 1
+        for k in [1, 5, 10]:
+            metrics[f"{args.datasetnames[0]}_R@{k}"] = np.mean(preds < k)
+        # map@10
+        metrics[f"{args.datasetnames[0]}_mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
+
+        logging.info(
+            f"Eval Epoch: {start_epoch} "
+            + "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in metrics.items()])
+        )
+
+        if args.wandb:
+            assert wandb is not None, "Please install wandb."
+            for name, val in metrics.items():
+                wandb.log({f"val/{name}": val, "epoch": start_epoch})
+
+
 if __name__ == '__main__':
     # (yusong) repeated run might have different metric results.
     # This is because we randomly select crop 10s for each audio.
     args = parse_args()
 
-    log_dir = os.path.dirname(os.path.dirname(args.resume))
+    if os.path.isdir(args.pretrained):
+        log_dir = os.path.dirname(args.pretrained)
+    else:
+        log_dir = os.path.dirname(os.path.dirname(args.pretrained))
+
     args.log_level = logging.DEBUG if args.debug else logging.INFO
     log_path = os.path.join(log_dir, 'out.log')
     setup_logging(log_path, args.log_level)
     params_file = os.path.join(log_dir, 'params.txt')
 
-    seed = 1234
+    seed = 3407
     random.seed(seed)
     torch.manual_seed(seed)
     torch.cuda.manual_seed(seed)
@@ -52,9 +119,10 @@ def find_params_value(file, key):
     cudnn.benchmark = True
     cudnn.deterministic = False
     pretrained = 'openai'
-    model_type = find_params_value(params_file, 'model')
+    amodel = find_params_value(params_file, 'amodel')
+    tmodel = find_params_value(params_file, 'tmodel')
 
-    if model_type is None:
+    if amodel is None or tmodel is None:
         raise ValueError('model type not found in params file')
 
     # set up dummy values for args
@@ -98,46 +166,74 @@ def find_params_value(file, key):
             proportion=1,
             dataset_path=args.datasetpath,
         )
+    model, model_cfg = create_model(
+        amodel,
+        tmodel,
+        pretrained,
+        precision='fp32',
+        device=device,
+        jit=False,
+        force_quick_gelu=False,
+        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+        skip_params=False,
+        enable_fusion=args.enable_fusion,
+        fusion_type=args.fusion_type
+    )  # a hack to get model_cfg
 
-    data = get_data(args, model_cfg=None)  # (yusong): hack: no model_cfg needed to get data
+    data = get_data(args, model_cfg=model_cfg)  # (yusong): hack: no model_cfg needed to get data
 
     writer = None  # if use tensorboard, initalize writer here
 
     if args.wandb:
         assert wandb is not None, "Please install wandb."
 
-        # find the line with "wandb_notes" and get the value
-        wandb_notes = find_params_value(params_file, 'wandb_notes')
-        if wandb_notes is None:
-            print(f'wandb_notes not found in params file: {params_file}, set to timestamp.')
-            wandb_notes = f'experiment_{time.strftime("%Y%m%d-%H%M%S")}'
-        wandb_notes = wandb_notes + '-eval-retrieval'
+        # # find the line with "wandb_notes" and get the value
+        # wandb_notes = find_params_value(params_file, 'wandb_notes')
+        # if wandb_notes is None:
+        #     print(f'wandb_notes not found in params file: {params_file}, set to timestamp.')
+        #     wandb_notes = f'experiment_{time.strftime("%Y%m%d-%H%M%S")}'
+        # wandb_notes = wandb_notes + '-eval-retrieval'
+        wandb_notes = args.wandb_notes
 
         logging.debug("Starting wandb.")
         args.train_sz = data["train"].dataloader.num_samples
         if args.val_data is not None:
             args.val_sz = data["val"].dataloader.num_samples
         # you will have to configure this for your project!
-        wandb.init(
-            project="clap",
-            notes=wandb_notes,
-            name=wandb_notes,
-            tags=[],
-            config=vars(args),
-        )
+        if args.wandb_id is not None:
+            wandb.init(
+                project="clap",
+                id=args.wandb_id,
+                resume=True
+            )
+        else:
+            wandb.init(
+                project="clap",
+                notes=wandb_notes,
+                name=wandb_notes,
+                tags=[],
+                config=vars(args),
+            )
         logging.debug("Finished loading wandb.")
 
-    all_model_checkpoints = sorted(glob.glob(os.path.join(log_dir, 'checkpoints', '*.pt')), key=os.path.getmtime)
+    if os.path.isdir(args.pretrained):
+        all_model_checkpoints = sorted(glob.glob(os.path.join(log_dir, 'checkpoints', '*.pt')), key=os.path.getmtime)
+    else:
+        all_model_checkpoints = [args.pretrained]
     for model_path in all_model_checkpoints:
         args.checkpoint_path = os.path.dirname(model_path)
         model, model_cfg = create_model(
-            model_type,
+            amodel,
+            tmodel,
             pretrained,
             precision='fp32',
             device=device,
             jit=False,
             force_quick_gelu=False,
-            openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir)
+            openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+            skip_params=False,
+            enable_fusion=args.enable_fusion,
+            fusion_type=args.fusion_type
         )
 
         # load model
@@ -164,4 +260,9 @@ def find_params_value(file, key):
         for param in model.parameters():
             param.requires_grad = False
 
-        evaluate(model, data, start_epoch, args, writer)
+        if args.datasetnames == ['esc50_no_overlap'] or args.datasetnames == ['VGGSound'] or \
+                args.datasetnames == ['usd8k_no_overlap']:
+            # use the same dataset for all models
+            evaluate_zeroshot(model, data, start_epoch, args, writer)
+        else:
+            evaluate(model, data, start_epoch, args, writer)
diff --git a/src/evaluate/gen_csv.py b/src/evaluate/gen_csv.py
index 4f3339f..4ecb2e3 100644
--- a/src/evaluate/gen_csv.py
+++ b/src/evaluate/gen_csv.py
@@ -14,8 +14,8 @@
 import os
 
 
-def get_output_from_single_audio(audio, text, model):
-    audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding")[
+def get_output_from_single_audio(audio, text, model, device):
+    audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding", device=device)[
         "embedding"
     ]
     if audio_embedding.ndim > 1:
@@ -38,8 +38,8 @@ def get_output_from_single_audio(audio, text, model):
     )
 
 
-def get_audio_embeddings(audio, model):
-    audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding")[
+def get_audio_embeddings(audio, model, device):
+    audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding", device=device)[
         "embedding"
     ]
     if audio_embedding.ndim > 1:
@@ -165,7 +165,7 @@ def get_text_embeddings(text, model):
                     audio_features,
                     audio_features_mlp,
                     logit_scale_a,
-                ) = get_audio_embeddings(audio, model)
+                ) = get_audio_embeddings(audio, model,device)
 
                 audio_features_all.append(audio_features)
                 audio_features_mlp_all.append(audio_features_mlp)
diff --git a/src/open_clip/factory.py b/src/open_clip/factory.py
index df881ba..3c3b286 100644
--- a/src/open_clip/factory.py
+++ b/src/open_clip/factory.py
@@ -75,7 +75,9 @@ def create_model(
     openai_model_cache_dir: str = os.path.expanduser("~/.cache/clip"),
     skip_params=True,
     pretrained_audio: str = "",
-    pretrained_text: str = ""
+    pretrained_text: str = "",
+    enable_fusion: bool = False,
+    fusion_type: str = 'None'
     # pretrained_image: bool = False,
 ):
     amodel_name = amodel_name.replace(
@@ -102,6 +104,8 @@ def create_model(
             device=device,
             jit=jit,
             cache_dir=openai_model_cache_dir,
+            enable_fusion=enable_fusion,
+            fusion_type=fusion_type
         )
         # See https://discuss.pytorch.org/t/valueerror-attemting-to-unscale-fp16-gradients/81372
         if precision == "amp" or precision == "fp32":
@@ -127,6 +131,8 @@ def create_model(
         #     else:
         #         assert False, 'pretrained image towers currently only supported for timm models'
         model_cfg["text_cfg"]["model_type"] = tmodel_name
+        model_cfg["enable_fusion"] = enable_fusion
+        model_cfg["fusion_type"] = fusion_type
         model = CLAP(**model_cfg)
 
         if pretrained:
@@ -153,20 +159,48 @@ def create_model(
 
         if pretrained_audio:
             if amodel_name.startswith('PANN'):
-                audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
-                audio_ckpt = audio_ckpt['model']
-                keys = list(audio_ckpt.keys())
-                for key in keys:
-                    v = audio_ckpt.pop(key)
-                    audio_ckpt['audio_branch.' + key] = v
+                if 'Cnn14_mAP' in pretrained_audio:  # official checkpoint
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                    audio_ckpt = audio_ckpt['model']
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if 'spectrogram_extractor' not in key and 'logmel_extractor' not in key:
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt['audio_branch.' + key] = v
+                elif os.path.basename(pretrained_audio).startswith('PANN'):  # checkpoint trained via HTSAT codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                    audio_ckpt = audio_ckpt['state_dict']
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if key.startswith('sed_model'):
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt['audio_branch.' + key[10:]] = v
+                elif os.path.basename(pretrained_audio).startswith('finetuned'):  # checkpoint trained via linear probe codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                else:
+                    raise ValueError('Unknown audio checkpoint')
             elif amodel_name.startswith('HTSAT'):
-                audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
-                audio_ckpt = audio_ckpt['state_dict']
-                keys = list(audio_ckpt.keys())
-                for key in keys:
-                    if key.startswith('sed_model'):
-                        v = audio_ckpt.pop(key)
-                        audio_ckpt['audio_branch.' + key[10:]] = v
+                if 'HTSAT_AudioSet_Saved' in pretrained_audio:  # official checkpoint
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                    audio_ckpt = audio_ckpt['state_dict']
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if key.startswith('sed_model') and ('spectrogram_extractor' not in key
+                                                            and 'logmel_extractor' not in key):
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt['audio_branch.' + key[10:]] = v
+                elif os.path.basename(pretrained_audio).startswith('HTSAT'):  # checkpoint trained via HTSAT codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                    audio_ckpt = audio_ckpt['state_dict']
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if key.startswith('sed_model'):
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt['audio_branch.' + key[10:]] = v
+                elif os.path.basename(pretrained_audio).startswith('finetuned'):  # checkpoint trained via linear probe codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                else:
+                    raise ValueError('Unknown audio checkpoint')
             else:
                 raise f'this audio encoder pretrained checkpoint is not support'
 
diff --git a/src/open_clip/feature_fusion.py b/src/open_clip/feature_fusion.py
new file mode 100644
index 0000000..c241951
--- /dev/null
+++ b/src/open_clip/feature_fusion.py
@@ -0,0 +1,193 @@
+'''
+Feature Fusion for Varible-Length Data Processing
+AFF/iAFF is referred and modified from https://github.com/YimianDai/open-aff/blob/master/aff_pytorch/aff_net/fusion.py
+According to the paper: Yimian Dai et al, Attentional Feature Fusion, IEEE Winter Conference on Applications of Computer Vision, WACV 2021
+'''
+
+import torch
+import torch.nn as nn
+
+
+class DAF(nn.Module):
+    '''
+    直接相加 DirectAddFuse
+    '''
+
+    def __init__(self):
+        super(DAF, self).__init__()
+
+    def forward(self, x, residual):
+        return x + residual
+
+
+class iAFF(nn.Module):
+    '''
+    多特征融合 iAFF
+    '''
+
+    def __init__(self, channels=64, r=4, type='2D'):
+        super(iAFF, self).__init__()
+        inter_channels = int(channels // r)
+
+        if type == '1D':
+            # 本地注意力
+            self.local_att = nn.Sequential(
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+
+            # 全局注意力
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool1d(1),
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+
+            # 第二次本地注意力
+            self.local_att2 = nn.Sequential(
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+            # 第二次全局注意力
+            self.global_att2 = nn.Sequential(
+                nn.AdaptiveAvgPool1d(1),
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+        elif type == '2D':
+            # 本地注意力
+            self.local_att = nn.Sequential(
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+
+            # 全局注意力
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool2d(1),
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+
+            # 第二次本地注意力
+            self.local_att2 = nn.Sequential(
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+            # 第二次全局注意力
+            self.global_att2 = nn.Sequential(
+                nn.AdaptiveAvgPool2d(1),
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+        else:
+            raise f'the type is not supported'
+
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x, residual):
+        flag = False
+        xa = x + residual
+        if xa.size(0) == 1:
+            xa = torch.cat([xa,xa],dim=0)
+            flag = True
+        xl = self.local_att(xa)
+        xg = self.global_att(xa)
+        xlg = xl + xg
+        wei = self.sigmoid(xlg)
+        xi = x * wei + residual * (1 - wei)
+
+        xl2 = self.local_att2(xi)
+        xg2 = self.global_att(xi)
+        xlg2 = xl2 + xg2
+        wei2 = self.sigmoid(xlg2)
+        xo = x * wei2 + residual * (1 - wei2)
+        if flag:
+            xo = xo[0].unsqueeze(0)
+        return xo
+
+
+class AFF(nn.Module):
+    '''
+    多特征融合 AFF
+    '''
+
+    def __init__(self, channels=64, r=4, type='2D'):
+        super(AFF, self).__init__()
+        inter_channels = int(channels // r)
+
+        if type == '1D':
+            self.local_att = nn.Sequential(
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool1d(1),
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+        elif type == '2D':
+            self.local_att = nn.Sequential(
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool2d(1),
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+        else:
+            raise f'the type is not supported.'
+        
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x, residual):
+        flag = False
+        xa = x + residual
+        if xa.size(0) == 1:
+            xa = torch.cat([xa,xa],dim=0)
+            flag = True
+        xl = self.local_att(xa)
+        xg = self.global_att(xa)
+        xlg = xl + xg
+        wei = self.sigmoid(xlg)
+        xo = 2 * x * wei + 2 * residual * (1 - wei)
+        if flag:
+            xo = xo[0].unsqueeze(0)
+        return xo
+
diff --git a/src/open_clip/htsat.py b/src/open_clip/htsat.py
index 3ed922e..db96116 100644
--- a/src/open_clip/htsat.py
+++ b/src/open_clip/htsat.py
@@ -24,6 +24,7 @@
 from itertools import repeat
 from .utils import do_mixup, interpolate
 
+from .feature_fusion import iAFF, AFF, DAF
 
 # from PyTorch internals
 def _ntuple(n):
@@ -70,7 +71,8 @@ def forward(self, x):
 class PatchEmbed(nn.Module):
     """ 2D Image to Patch Embedding
     """
-    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True, patch_stride = 16):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True, patch_stride = 16,
+        enable_fusion=False, fusion_type='None'):
         super().__init__()
         img_size = to_2tuple(img_size)
         patch_size = to_2tuple(patch_size)
@@ -83,17 +85,59 @@ def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_
         self.flatten = flatten
         self.in_chans = in_chans
         self.embed_dim = embed_dim
+
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
         
         padding = ((patch_size[0] - patch_stride[0]) // 2, (patch_size[1] - patch_stride[1]) // 2)
 
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_stride, padding=padding)
+        if (self.enable_fusion) and (self.fusion_type == 'channel_map'):
+            self.proj = nn.Conv2d(in_chans*4, embed_dim, kernel_size=patch_size, stride=patch_stride, padding=padding)
+        else:
+            self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_stride, padding=padding)
         self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
 
-    def forward(self, x):
-        B, C, H, W = x.shape
-        assert H == self.img_size[0] and W == self.img_size[1], \
-            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
-        x = self.proj(x)
+        if (self.enable_fusion) and (self.fusion_type in ['daf_2d','aff_2d','iaff_2d']):
+            self.mel_conv2d = nn.Conv2d(in_chans, embed_dim, kernel_size=(patch_size[0], patch_size[1]*3), stride=(patch_stride[0], patch_stride[1] * 3), padding=padding)
+            if self.fusion_type == 'daf_2d':
+                self.fusion_model = DAF()
+            elif self.fusion_type == 'aff_2d':
+                self.fusion_model = AFF(channels=embed_dim, type='2D')
+            elif self.fusion_type == 'iaff_2d':
+                self.fusion_model = iAFF(channels=embed_dim, type='2D')    
+    def forward(self, x, longer_idx = None):
+        if (self.enable_fusion) and (self.fusion_type in ['daf_2d','aff_2d','iaff_2d']):
+            global_x = x[:,0:1,:,:]
+            
+
+            # global processing
+            B, C, H, W = global_x.shape
+            assert H == self.img_size[0] and W == self.img_size[1], \
+                f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+            global_x = self.proj(global_x)
+            TW = global_x.size(-1)
+            if len(longer_idx) > 0:
+                # local processing
+                local_x = x[longer_idx,1:,:,:].contiguous()
+                B, C, H, W = local_x.shape
+                local_x = local_x.view(B*C,1,H,W)
+                local_x = self.mel_conv2d(local_x)
+                local_x = local_x.view(B,C,local_x.size(1),local_x.size(2),local_x.size(3))
+                local_x = local_x.permute((0,2,3,1,4)).contiguous().flatten(3)
+                TB,TC,TH,_ = local_x.size()
+                if local_x.size(-1) < TW:
+                    local_x = torch.cat([local_x, torch.zeros((TB,TC,TH,TW-local_x.size(-1)), device=global_x.device)], dim=-1)
+                else:
+                    local_x = local_x[:,:,:,:TW]
+                
+                global_x[longer_idx] = self.fusion_model(global_x[longer_idx],local_x)
+            x = global_x
+        else:
+            B, C, H, W = x.shape
+            assert H == self.img_size[0] and W == self.img_size[1], \
+                f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+            x = self.proj(x)
+        
         if self.flatten:
             x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
         x = self.norm(x)
@@ -584,7 +628,8 @@ def __init__(self, spec_size=256, patch_size=4, patch_stride=(4,4),
                  drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
                  norm_layer=nn.LayerNorm, 
                  ape=False, patch_norm=True,
-                 use_checkpoint=False, norm_before_mlp='ln', config = None, **kwargs):
+                 use_checkpoint=False, norm_before_mlp='ln', config = None, 
+                 enable_fusion = False, fusion_type = 'None', **kwargs):
         super(HTSAT_Swin_Transformer, self).__init__()
 
         self.config = config
@@ -615,6 +660,9 @@ def __init__(self, spec_size=256, patch_size=4, patch_stride=(4,4),
 
         self.use_checkpoint = use_checkpoint
 
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+
         #  process mel-spec ; used only once
         self.freq_ratio = self.spec_size // self.config.mel_bins
         window = 'hann'
@@ -641,7 +689,9 @@ def __init__(self, spec_size=256, patch_size=4, patch_stride=(4,4),
         # split spctrogram into non-overlapping patches
         self.patch_embed = PatchEmbed(
             img_size=self.spec_size, patch_size=self.patch_size, in_chans=self.in_chans, 
-            embed_dim=self.embed_dim, norm_layer=self.norm_layer, patch_stride = patch_stride)
+            embed_dim=self.embed_dim, norm_layer=self.norm_layer, patch_stride = patch_stride,
+            enable_fusion=self.enable_fusion, fusion_type=self.fusion_type
+            )
 
         num_patches = self.patch_embed.num_patches
         patches_resolution = self.patch_embed.grid_size
@@ -688,6 +738,19 @@ def __init__(self, spec_size=256, patch_size=4, patch_stride=(4,4),
             padding = (0,1)
         )
         self.head = nn.Linear(num_classes, num_classes)
+
+        if (self.enable_fusion) and (self.fusion_type in ['daf_1d','aff_1d','iaff_1d']):
+            self.mel_conv1d = nn.Sequential(
+                nn.Conv1d(64, 64, kernel_size=5, stride=3, padding=2),
+                nn.BatchNorm1d(64)
+            )
+            if self.fusion_type == 'daf_1d':
+                self.fusion_model = DAF()
+            elif self.fusion_type == 'aff_1d':
+                self.fusion_model = AFF(channels=64, type='1D')
+            elif self.fusion_type == 'iaff_1d':
+                self.fusion_model = iAFF(channels=64, type='1D')
+                
         self.apply(self._init_weights)
 
     def _init_weights(self, m):
@@ -708,11 +771,11 @@ def no_weight_decay_keywords(self):
         return {'relative_position_bias_table'}
 
 
-    def forward_features(self, x):
+    def forward_features(self, x, longer_idx = None):
         # A deprecated optimization for using a hierarchical output from different blocks
 
         frames_num = x.shape[2]        
-        x = self.patch_embed(x)
+        x = self.patch_embed(x, longer_idx = longer_idx)
         if self.ape:
             x = x + self.absolute_pos_embed
         x = self.pos_drop(x)
@@ -800,59 +863,113 @@ def repeat_wat2img(self, x, cur_pos):
         x = x.repeat(repeats = (1,1,4,1))
         return x
 
-    def forward(self, x: torch.Tensor, mixup_lambda = None, infer_mode = False):# out_feat_keys: List[str] = None):
-        x = self.spectrogram_extractor(x)   # (batch_size, 1, time_steps, freq_bins)
-        x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
-
-        x = x.transpose(1, 3)
-        x = self.bn0(x)
-        x = x.transpose(1, 3)
-        if self.training:
-            x = self.spec_augmenter(x)
-        if self.training and mixup_lambda is not None:
-            x = do_mixup(x, mixup_lambda)
-        
-        if infer_mode:
-            # in infer mode. we need to handle different length audio input
-            frame_num = x.shape[2]
-            target_T = int(self.spec_size * self.freq_ratio)
-            repeat_ratio = math.floor(target_T / frame_num)
-            x = x.repeat(repeats=(1,1,repeat_ratio,1))
+    def forward(self, x: torch.Tensor, mixup_lambda = None, infer_mode = False, device=None):# out_feat_keys: List[str] = None):
+
+        if self.enable_fusion and x["longer"].sum() == 0:
+            # if no audio is longer than 10s, then randomly select one audio to be longer
+            x["longer"][torch.randint(0, x["longer"].shape[0], (1,))] = True
+
+        if not self.enable_fusion:
+            x = x["waveform"].to(device=device, non_blocking=True)
+            x = self.spectrogram_extractor(x)   # (batch_size, 1, time_steps, freq_bins)
+            x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
+            x = x.transpose(1, 3)
+            x = self.bn0(x)
+            x = x.transpose(1, 3)
+            if self.training:
+                x = self.spec_augmenter(x)
+
+            if self.training and mixup_lambda is not None:
+                x = do_mixup(x, mixup_lambda)
+                
             x = self.reshape_wav2img(x)
             output_dict = self.forward_features(x)
         else:
-            if x.shape[2] > self.freq_ratio * self.spec_size:
-                if self.training:
-                    x = self.crop_wav(x, crop_size=self.freq_ratio * self.spec_size)
-                    x = self.reshape_wav2img(x)
-                    output_dict = self.forward_features(x)
+            longer_list = x["longer"].to(device=device, non_blocking=True)
+            x = x["mel_fusion"].to(device=device, non_blocking=True)
+            x = x.transpose(1, 3)
+            x = self.bn0(x)
+            x = x.transpose(1, 3)
+            longer_list_idx = torch.where(longer_list)[0]
+            if self.fusion_type in ['daf_1d','aff_1d','iaff_1d']:
+                new_x = x[:,0:1,:,:].clone().contiguous()
+                if len(longer_list_idx) > 0:
+                # local processing
+                    fusion_x_local = x[longer_list_idx,1:,:,:].clone().contiguous()
+                    FB,FC,FT,FF = fusion_x_local.size()
+                    fusion_x_local = fusion_x_local.view(FB * FC, FT, FF)
+                    fusion_x_local = torch.permute(fusion_x_local, (0,2,1)).contiguous()
+                    fusion_x_local = self.mel_conv1d(fusion_x_local)
+                    fusion_x_local = fusion_x_local.view(FB,FC,FF,fusion_x_local.size(-1))
+                    fusion_x_local = torch.permute(fusion_x_local, (0,2,1,3)).contiguous().flatten(2)
+                    if fusion_x_local.size(-1) < FT:
+                        fusion_x_local = torch.cat([fusion_x_local, torch.zeros((FB,FF,FT- fusion_x_local.size(-1)), device=device)], dim=-1)
+                    else:
+                        fusion_x_local = fusion_x_local[:,:,:FT]
+                    # 1D fusion
+                    new_x = new_x.squeeze(1).permute((0,2,1)).contiguous()
+                    new_x[longer_list_idx] = self.fusion_model(new_x[longer_list_idx], fusion_x_local)
+                    x = new_x.permute((0,2,1)).contiguous()[:,None,:,:]
                 else:
-                    # Change: Hard code here
-                    overlap_size = (x.shape[2] - 1) // 4
-                    output_dicts = []
-                    crop_size = (x.shape[2] - 1) // 2
-                    for cur_pos in range(0, x.shape[2] - crop_size - 1, overlap_size):
-                        tx = self.crop_wav(x, crop_size = crop_size, spe_pos = cur_pos)
-                        tx = self.reshape_wav2img(tx)
-                        output_dicts.append(self.forward_features(tx))
-                    clipwise_output = torch.zeros_like(output_dicts[0]["clipwise_output"]).float().to(x.device)
-                    framewise_output = torch.zeros_like(output_dicts[0]["framewise_output"]).float().to(x.device)
-                    for d in output_dicts:
-                        clipwise_output += d["clipwise_output"]
-                        framewise_output += d["framewise_output"]
-                    clipwise_output  = clipwise_output / len(output_dicts)
-                    framewise_output = framewise_output / len(output_dicts)
-                    output_dict = {
-                        'framewise_output': framewise_output, 
-                        'clipwise_output': clipwise_output
-                    }
-            else: # this part is typically used, and most easy one
-                x = self.reshape_wav2img(x)
-                output_dict = self.forward_features(x)
+                    x = new_x
+
+            elif self.fusion_type in ['daf_2d','aff_2d','iaff_2d','channel_map']:
+                x = x # no change
+
+            if self.training:
+                x = self.spec_augmenter(x)
+            if self.training and mixup_lambda is not None:
+                x = do_mixup(x, mixup_lambda)
+
+            x = self.reshape_wav2img(x)
+            output_dict = self.forward_features(x, longer_idx = longer_list_idx)
+       
+        # if infer_mode:
+        #     # in infer mode. we need to handle different length audio input
+        #     frame_num = x.shape[2]
+        #     target_T = int(self.spec_size * self.freq_ratio)
+        #     repeat_ratio = math.floor(target_T / frame_num)
+        #     x = x.repeat(repeats=(1,1,repeat_ratio,1))
+        #     x = self.reshape_wav2img(x)
+        #     output_dict = self.forward_features(x)
+        # else:
+        #     if x.shape[2] > self.freq_ratio * self.spec_size:
+        #         if self.training:
+        #             x = self.crop_wav(x, crop_size=self.freq_ratio * self.spec_size)
+        #             x = self.reshape_wav2img(x)
+        #             output_dict = self.forward_features(x)
+        #         else:
+        #             # Change: Hard code here
+        #             overlap_size = (x.shape[2] - 1) // 4
+        #             output_dicts = []
+        #             crop_size = (x.shape[2] - 1) // 2
+        #             for cur_pos in range(0, x.shape[2] - crop_size - 1, overlap_size):
+        #                 tx = self.crop_wav(x, crop_size = crop_size, spe_pos = cur_pos)
+        #                 tx = self.reshape_wav2img(tx)
+        #                 output_dicts.append(self.forward_features(tx))
+        #             clipwise_output = torch.zeros_like(output_dicts[0]["clipwise_output"]).float().to(x.device)
+        #             framewise_output = torch.zeros_like(output_dicts[0]["framewise_output"]).float().to(x.device)
+        #             for d in output_dicts:
+        #                 clipwise_output += d["clipwise_output"]
+        #                 framewise_output += d["framewise_output"]
+        #             clipwise_output  = clipwise_output / len(output_dicts)
+        #             framewise_output = framewise_output / len(output_dicts)
+        #             output_dict = {
+        #                 'framewise_output': framewise_output, 
+        #                 'clipwise_output': clipwise_output
+        #             }
+        #     else: # this part is typically used, and most easy one
+        #         x = self.reshape_wav2img(x)
+        #         output_dict = self.forward_features(x)
         # x = self.head(x)
+
+        # We process the data in the dataloader part, in that here we only consider the input_T < fixed_T
+        
+        
+
         return output_dict
 
-def create_htsat_model(audio_cfg):
+def create_htsat_model(audio_cfg, enable_fusion=False, fusion_type='None'):
     try:
 
         assert audio_cfg.model_name in ["tiny", "base", "large"], "model name for HTS-AT is wrong!"
@@ -866,7 +983,9 @@ def create_htsat_model(audio_cfg):
                 depths=[2,2,6,2],
                 num_heads=[4,8,16,32],
                 window_size=8,
-                config = audio_cfg
+                config = audio_cfg,
+                enable_fusion = enable_fusion,
+                fusion_type = fusion_type
             )
         elif audio_cfg.model_name == "base":
             model = HTSAT_Swin_Transformer(
@@ -878,7 +997,9 @@ def create_htsat_model(audio_cfg):
                 depths=[2,2,12,2],
                 num_heads=[4,8,16,32],
                 window_size=8,
-                config = audio_cfg
+                config = audio_cfg,
+                enable_fusion = enable_fusion,
+                fusion_type = fusion_type
             )
         elif audio_cfg.model_name == "large":
             model = HTSAT_Swin_Transformer(
@@ -890,7 +1011,9 @@ def create_htsat_model(audio_cfg):
                 depths=[2,2,12,2],
                 num_heads=[4,8,16,32],
                 window_size=8,
-                config = audio_cfg
+                config = audio_cfg,
+                enable_fusion = enable_fusion,
+                fusion_type = fusion_type
             )
         
         return model
diff --git a/src/open_clip/linear_probe.py b/src/open_clip/linear_probe.py
index 3dc92a1..bb2841d 100644
--- a/src/open_clip/linear_probe.py
+++ b/src/open_clip/linear_probe.py
@@ -3,8 +3,9 @@
 from torch import nn
 from .model import MLPLayers
 
+
 class LinearProbe(nn.Module):
-    def __init__(self, model, mlp, freeze, in_ch, out_ch, act = None):
+    def __init__(self, model, mlp, freeze, in_ch, out_ch, act=None):
         """
         Args:
             model: nn.Module
@@ -17,16 +18,17 @@ def __init__(self, model, mlp, freeze, in_ch, out_ch, act = None):
         super().__init__()
         in_ch = 512
         self.clap_model = model
+        self.clap_model.text_branch = None  # to save memory
         self.freeze = freeze
         if mlp:
-            self.lp_layer = MLPLayers(units = [in_ch, in_ch * 2, out_ch])
+            self.lp_layer = MLPLayers(units=[in_ch, in_ch * 2, out_ch])
         else:
             self.lp_layer = nn.Linear(in_ch, out_ch)
-        
+
         if self.freeze:
             for param in self.clap_model.parameters():
                 param.requires_grad = False
-        
+
         if act == 'None':
             self.act = None
         elif act == 'relu':
@@ -40,11 +42,11 @@ def __init__(self, model, mlp, freeze, in_ch, out_ch, act = None):
         elif act == 'sigmoid':
             self.act = nn.Sigmoid()
 
-    def forward(self, x):
+    def forward(self, x, mix_lambda=None, device=None):
         """
         Args:
-            x: waveform, torch.tensor [batch, t_samples]
-                
+            x: waveform, torch.tensor [batch, t_samples] / batch of mel_spec and longer list
+            mix_lambda: torch.tensor [batch], the mixup lambda
         Returns:
             class_prob: torch.tensor [batch, class_num]
 
@@ -53,9 +55,9 @@ def forward(self, x):
         if self.freeze:
             self.clap_model.eval()
 
-        x = self.clap_model(audio=x, text=None)
+        x = self.clap_model.audio_projection(
+            self.clap_model.audio_branch(x, mixup_lambda=mix_lambda, device=device)["embedding"])
         out = self.lp_layer(x)
         if self.act is not None:
             out = self.act(out)
         return out
-
diff --git a/src/open_clip/model.py b/src/open_clip/model.py
index f3ade3b..7af6507 100644
--- a/src/open_clip/model.py
+++ b/src/open_clip/model.py
@@ -426,6 +426,10 @@ def __init__(
         audio_cfg: CLAPAudioCfp,
         text_cfg: CLAPTextCfg,
         quick_gelu: bool = False,
+        enable_fusion: bool = False,
+        fusion_type: str = 'None',
+        joint_embed_shape: int = 512,
+        mlp_act: str = 'relu',
     ):
         super().__init__()
         if isinstance(audio_cfg, dict):
@@ -435,6 +439,11 @@ def __init__(
 
         self.audio_cfg = audio_cfg
         self.text_cfg = text_cfg
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+        self.joint_embed_shape = joint_embed_shape
+        self.mlp_act = mlp_act
+
 
         self.context_length = text_cfg.context_length
 
@@ -443,12 +452,19 @@ def __init__(
         # NOTE: timm models always use native GELU regardless of quick_gelu flag.
         act_layer = QuickGELU if quick_gelu else nn.GELU
 
+        if mlp_act == 'relu':
+            mlp_act_layer = nn.ReLU()
+        elif mlp_act == 'gelu':
+            mlp_act_layer = nn.GELU()
+        else:
+            raise NotImplementedError
+
         # audio branch
         # audio branch parameters
         if audio_cfg.model_type == "PANN":
-            self.audio_branch = create_pann_model(audio_cfg)
+            self.audio_branch = create_pann_model(audio_cfg, enable_fusion, fusion_type)
         elif audio_cfg.model_type == "HTSAT":
-            self.audio_branch = create_htsat_model(audio_cfg)
+            self.audio_branch = create_htsat_model(audio_cfg, enable_fusion, fusion_type)
         else:
             logging.error(f"Model config for {audio_cfg.model_type} not found")
             raise RuntimeError(f"Model config for {audio_cfg.model_type} not found.")
@@ -468,28 +484,44 @@ def __init__(
                 torch.empty(self.context_length, text_cfg.width)
             )
             self.ln_final = LayerNorm(text_cfg.width)
-            self.text_transform = MLPLayers(units=[512, 512, 512], dropout=0.1)
-            self.text_projection = nn.Parameter(
-                torch.empty(text_cfg.width, 512)
-            )  # HARDCORE as 512
+            self.text_transform = MLPLayers(units=[self.joint_embed_shape,
+                                                   self.joint_embed_shape,
+                                                   self.joint_embed_shape], dropout=0.1)
+            self.text_projection = nn.Sequential(
+                nn.Linear(text_cfg.width, self.joint_embed_shape),
+                mlp_act_layer,
+                nn.Linear(self.joint_embed_shape, self.joint_embed_shape)
+            )
         elif text_cfg.model_type == "bert":
             self.text_branch = BertModel.from_pretrained("bert-base-uncased")
-            self.text_transform = MLPLayers(units=[512, 512, 512], dropout=0.1)
-            self.text_projection = nn.Parameter(
-                torch.empty(768, 512)
-            )  # HARDCORE as 512
+            self.text_transform = MLPLayers(units=[self.joint_embed_shape,
+                                                   self.joint_embed_shape,
+                                                   self.joint_embed_shape], dropout=0.1)
+            self.text_projection = nn.Sequential(
+                nn.Linear(768, self.joint_embed_shape),
+                mlp_act_layer,
+                nn.Linear(self.joint_embed_shape, self.joint_embed_shape)
+            )
         elif text_cfg.model_type == "roberta":
             self.text_branch = RobertaModel.from_pretrained('roberta-base')
-            self.text_transform = MLPLayers(units=[512, 512, 512], dropout=0.1)
-            self.text_projection = nn.Parameter(
-                torch.empty(768, 512)
-            )  # HARDCORE as 512
+            self.text_transform = MLPLayers(units=[self.joint_embed_shape,
+                                                   self.joint_embed_shape,
+                                                   self.joint_embed_shape], dropout=0.1)
+            self.text_projection = nn.Sequential(
+                nn.Linear(768, self.joint_embed_shape),
+                mlp_act_layer,
+                nn.Linear(self.joint_embed_shape, self.joint_embed_shape)
+            )
         elif text_cfg.model_type == "bart":
             self.text_branch = BartModel.from_pretrained('facebook/bart-base')
-            self.text_transform = MLPLayers(units=[512, 512, 512], dropout=0.1)
-            self.text_projection = nn.Parameter(
-                torch.empty(768, 512)
-            )  # HARDCORE as 512
+            self.text_transform = MLPLayers(units=[self.joint_embed_shape,
+                                                   self.joint_embed_shape,
+                                                   self.joint_embed_shape], dropout=0.1)
+            self.text_projection = nn.Sequential(
+                nn.Linear(768, self.joint_embed_shape),
+                mlp_act_layer,
+                nn.Linear(self.joint_embed_shape, self.joint_embed_shape)
+            )
         else:
             logging.error(f"Model config for {text_cfg.model_type} not found")
             raise RuntimeError(f"Model config for {text_cfg.model_type} not found.")
@@ -497,12 +529,18 @@ def __init__(
         # text branch parameters
 
         # audio branch parameters
-        self.audio_transform = MLPLayers(units=[512, 512, 512], dropout=0.1)
+        self.audio_transform = MLPLayers(units=[self.joint_embed_shape,
+                                                self.joint_embed_shape,
+                                                self.joint_embed_shape], dropout=0.1)
 
         # below here is text branch parameters
 
         # ============================================================================================================
-        self.audio_projection = MLPLayers(units=[embed_dim, 512, 512], dropout=0.1)
+        self.audio_projection = nn.Sequential(
+                nn.Linear(embed_dim, self.joint_embed_shape),
+                mlp_act_layer,
+                nn.Linear(self.joint_embed_shape, self.joint_embed_shape)
+            )
 
         self.logit_scale_a = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
         self.logit_scale_t = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
@@ -537,8 +575,8 @@ def init_text_branch_parameters(self):
         # if hasattr(self.visual, 'init_parameters'):
         # self.visual.init_parameters()
 
-        if self.text_projection is not None:
-            nn.init.normal_(self.text_projection, std=width**-0.5)
+        # if self.text_projection is not None:
+        #     nn.init.normal_(self.text_projection, std=width**-0.5)
 
     def build_attention_mask(self):
         # lazily create causal attention mask, with full attention between the vision tokens
@@ -548,8 +586,8 @@ def build_attention_mask(self):
         mask.triu_(1)  # zero out the lower diagonal
         return mask
 
-    def encode_audio(self, audio):
-        return self.audio_branch(audio, None)  # mix lambda needs to add
+    def encode_audio(self, audio, device):
+        return self.audio_branch(audio, mixup_lambda=None, device=device)  # mix lambda needs to add
 
     # def list_of_dict_of_tensor2dict_of_tensor(self, x, device):
     #     tmp = {}
@@ -574,7 +612,7 @@ def encode_text(self, text, device):
 
             # x.shape = [batch_size, n_ctx, transformer.width]
             # take features from the eot embedding (eot_token is the highest number in each sequence)
-            x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+            x = self.text_projection(x[torch.arange(x.shape[0]), text.argmax(dim=-1)])
         elif self.text_branch_type == "bert":
             # text = self.list_of_dict_of_tensor2dict_of_tensor(text, device)
             # text = BatchEncoding(text)
@@ -587,7 +625,7 @@ def encode_text(self, text, device):
                     device=device, non_blocking=True
                 ),
             )["pooler_output"]
-            x = x @ self.text_projection
+            x = self.text_projection(x)
         elif self.text_branch_type == "roberta":
             x = self.text_branch(
                 input_ids=text["input_ids"].to(device=device, non_blocking=True),
@@ -595,7 +633,7 @@ def encode_text(self, text, device):
                     device=device, non_blocking=True
                 ),
             )["pooler_output"]
-            x = x @ self.text_projection
+            x = self.text_projection(x)
         elif self.text_branch_type == "bart":
             x = torch.mean(self.text_branch(
                 input_ids=text["input_ids"].to(device=device, non_blocking=True),
@@ -603,7 +641,7 @@ def encode_text(self, text, device):
                     device=device, non_blocking=True
                 ),
             )["encoder_last_hidden_state"],axis=1)
-            x = x @ self.text_projection
+            x = self.text_projection(x)
         else:
             logging.error(f"Model type {self.text_branch_type} not found")
             raise RuntimeError(f"Model type {self.text_branch_type} not found.")
@@ -615,7 +653,7 @@ def forward(self, audio, text, device=None):
         Parameters
         ----------
         audio: torch.Tensor (batch_size, audio_length)
-            the time-domain audio input
+            the time-domain audio input / the batch of mel_spec and longer list.
         text: torch.Tensor () // need to add
             the text token input
         """
@@ -624,11 +662,14 @@ def forward(self, audio, text, device=None):
                 device = audio.device
             elif text is not None:
                 device = text.device
-        if audio is None:
+        if audio is None and text is None:
+            # a hack to get the logit scale
+            return self.logit_scale_a.exp(), self.logit_scale_t.exp()
+        elif audio is None:
             return self.encode_text(text, device=device)
         elif text is None:
-            return self.audio_projection(self.encode_audio(audio)["embedding"])
-        audio_features = self.audio_projection(self.encode_audio(audio)["embedding"])
+            return self.audio_projection(self.encode_audio(audio, device=device)["embedding"])
+        audio_features = self.audio_projection(self.encode_audio(audio, device=device)["embedding"])
         audio_features = F.normalize(audio_features, dim=-1)
 
         text_features = self.encode_text(
@@ -651,7 +692,10 @@ def forward(self, audio, text, device=None):
             self.logit_scale_t.exp(),
         )
 
-    def audio_infer(self, audio, hopsize=None, key="embedding"):
+    def get_logit_scale(self):
+        return self.logit_scale_a.exp(), self.logit_scale_t.exp()
+
+    def audio_infer(self, audio, hopsize=None, key="embedding", device=None):
         """Forward one audio and produce the audio embedding
 
         Parameters
@@ -679,7 +723,7 @@ def audio_infer(self, audio, hopsize=None, key="embedding"):
         # PANN
         if self.audio_cfg.model_type == "PANN":
             audio_input = audio.unsqueeze(dim=0)
-            output_dict[key] = self.encode_audio(audio_input)[key].squeeze(dim=0)
+            output_dict[key] = self.encode_audio(audio_input, device=device)[key].squeeze(dim=0)
         elif self.audio_cfg.model_type == "HTSAT":
             # repeat
             audio_len = len(audio)
@@ -700,10 +744,10 @@ def audio_infer(self, audio, hopsize=None, key="embedding"):
                 ]
                 audio_input.append(audio[-self.audio_cfg.clip_samples :].clone())
                 audio_input = torch.stack(audio_input)
-                output_dict[key] = self.encode_audio(audio_input)[key]
+                output_dict[key] = self.encode_audio(audio_input, device=device)[key]
             else:
                 audio_input = audio.unsqueeze(dim=0)
-                output_dict[key] = self.encode_audio(audio_input)[key].squeeze(dim=0)
+                output_dict[key] = self.encode_audio(audio_input, device=device)[key].squeeze(dim=0)
 
         return output_dict
 
@@ -738,7 +782,7 @@ def _convert_weights_to_fp16(l):
 
 
 # Ignore the state dict of the vision part
-def build_model_from_openai_state_dict(state_dict: dict, model_cfg):
+def build_model_from_openai_state_dict(state_dict: dict, model_cfg, enable_fusion: bool = False, fusion_type: str = 'None'):
 
     embed_dim = model_cfg["embed_dim"]
     audio_cfg = model_cfg["audio_cfg"]
@@ -763,6 +807,8 @@ def build_model_from_openai_state_dict(state_dict: dict, model_cfg):
         audio_cfg=audio_cfg,
         text_cfg=text_cfg,
         quick_gelu=True,  # OpenAI models were trained with QuickGELU
+        enable_fusion=enable_fusion,
+        fusion_type=fusion_type
     )
     state_dict["logit_scale_a"] = state_dict["logit_scale"]
     state_dict["logit_scale_t"] = state_dict["logit_scale"]
diff --git a/src/open_clip/model_configs/HTSAT-tiny-win-1536.json b/src/open_clip/model_configs/HTSAT-tiny-win-1536.json
new file mode 100644
index 0000000..73e4299
--- /dev/null
+++ b/src/open_clip/model_configs/HTSAT-tiny-win-1536.json
@@ -0,0 +1,23 @@
+{
+    "embed_dim": 768,
+    "audio_cfg": {
+        "audio_length": 1024,
+        "clip_samples": 480000,
+        "mel_bins": 64,
+        "sample_rate": 48000,
+        "window_size": 1536,
+        "hop_size": 480,
+        "fmin": 50,
+        "fmax": 14000,
+        "class_num": 527,
+        "model_type": "HTSAT",
+        "model_name": "tiny"
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}
\ No newline at end of file
diff --git a/src/open_clip/model_configs/PANN-14-fmax-18k.json b/src/open_clip/model_configs/PANN-14-fmax-18k.json
new file mode 100644
index 0000000..b7989bc
--- /dev/null
+++ b/src/open_clip/model_configs/PANN-14-fmax-18k.json
@@ -0,0 +1,23 @@
+{
+    "embed_dim": 2048,
+    "audio_cfg": {
+        "audio_length": 1024,
+        "clip_samples": 480000,
+        "mel_bins": 64,
+        "sample_rate": 48000,
+        "window_size": 1024,
+        "hop_size": 480,
+        "fmin": 50,
+        "fmax": 18000,
+        "class_num": 527,
+        "model_type": "PANN",
+        "model_name": "Cnn14"
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}
\ No newline at end of file
diff --git a/src/open_clip/model_configs/PANN-14-fmax-8k-20s.json b/src/open_clip/model_configs/PANN-14-fmax-8k-20s.json
new file mode 100644
index 0000000..56bdb56
--- /dev/null
+++ b/src/open_clip/model_configs/PANN-14-fmax-8k-20s.json
@@ -0,0 +1,23 @@
+{
+    "embed_dim": 2048,
+    "audio_cfg": {
+        "audio_length": 1024,
+        "clip_samples": 960000,
+        "mel_bins": 64,
+        "sample_rate": 48000,
+        "window_size": 1024,
+        "hop_size": 360,
+        "fmin": 50,
+        "fmax": 8000,
+        "class_num": 527,
+        "model_type": "PANN",
+        "model_name": "Cnn14"
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}
\ No newline at end of file
diff --git a/src/open_clip/model_configs/PANN-14-win-1536.json b/src/open_clip/model_configs/PANN-14-win-1536.json
new file mode 100644
index 0000000..5a9e7e2
--- /dev/null
+++ b/src/open_clip/model_configs/PANN-14-win-1536.json
@@ -0,0 +1,23 @@
+{
+    "embed_dim": 2048,
+    "audio_cfg": {
+        "audio_length": 1024,
+        "clip_samples": 480000,
+        "mel_bins": 64,
+        "sample_rate": 48000,
+        "window_size": 1536,
+        "hop_size": 480,
+        "fmin": 50,
+        "fmax": 14000,
+        "class_num": 527,
+        "model_type": "PANN",
+        "model_name": "Cnn14"
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}
\ No newline at end of file
diff --git a/src/open_clip/openai.py b/src/open_clip/openai.py
index 4ad284d..9911b6e 100644
--- a/src/open_clip/openai.py
+++ b/src/open_clip/openai.py
@@ -26,6 +26,8 @@ def load_openai_model(
         device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu",
         jit=True,
         cache_dir=os.path.expanduser("~/.cache/clip"),
+        enable_fusion: bool = False,
+        fusion_type: str = 'None'
 ):
     """Load a CLIP model, preserve its text pretrained part, and set in the CLAP model
 
@@ -65,10 +67,10 @@ def load_openai_model(
 
     if not jit:
         try:
-            model = build_model_from_openai_state_dict(state_dict or model.state_dict(), model_cfg).to(device)
+            model = build_model_from_openai_state_dict(state_dict or model.state_dict(), model_cfg, enable_fusion, fusion_type).to(device)
         except KeyError:
             sd = {k[7:]: v for k, v in state_dict["state_dict"].items()}
-            model = build_model_from_openai_state_dict(sd, model_cfg).to(device)
+            model = build_model_from_openai_state_dict(sd, model_cfg, enable_fusion, fusion_type).to(device)
 
         if str(device) == "cpu":
             model.float()
diff --git a/src/open_clip/pann_model.py b/src/open_clip/pann_model.py
index ad5efc5..109db5f 100644
--- a/src/open_clip/pann_model.py
+++ b/src/open_clip/pann_model.py
@@ -11,17 +11,18 @@
 from torchlibrosa.augmentation import SpecAugmentation
 
 from .utils import do_mixup, interpolate, pad_framewise_output
- 
+from .feature_fusion import iAFF, AFF, DAF
+
 
 def init_layer(layer):
     """Initialize a Linear or Convolutional layer. """
     nn.init.xavier_uniform_(layer.weight)
- 
+
     if hasattr(layer, 'bias'):
         if layer.bias is not None:
             layer.bias.data.fill_(0.)
-            
-    
+
+
 def init_bn(bn):
     """Initialize a Batchnorm layer. """
     bn.bias.data.fill_(0.)
@@ -30,33 +31,33 @@ def init_bn(bn):
 
 class ConvBlock(nn.Module):
     def __init__(self, in_channels, out_channels):
-        
+
         super(ConvBlock, self).__init__()
-        
-        self.conv1 = nn.Conv2d(in_channels=in_channels, 
+
+        self.conv1 = nn.Conv2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=(3, 3), stride=(1, 1),
                               padding=(1, 1), bias=False)
-                              
-        self.conv2 = nn.Conv2d(in_channels=out_channels, 
+
+        self.conv2 = nn.Conv2d(in_channels=out_channels,
                               out_channels=out_channels,
                               kernel_size=(3, 3), stride=(1, 1),
                               padding=(1, 1), bias=False)
-                              
+
         self.bn1 = nn.BatchNorm2d(out_channels)
         self.bn2 = nn.BatchNorm2d(out_channels)
 
         self.init_weight()
-        
+
     def init_weight(self):
         init_layer(self.conv1)
         init_layer(self.conv2)
         init_bn(self.bn1)
         init_bn(self.bn2)
 
-        
+
     def forward(self, input, pool_size=(2, 2), pool_type='avg'):
-        
+
         x = input
         x = F.relu_(self.bn1(self.conv1(x)))
         x = F.relu_(self.bn2(self.conv2(x)))
@@ -70,31 +71,31 @@ def forward(self, input, pool_size=(2, 2), pool_type='avg'):
             x = x1 + x2
         else:
             raise Exception('Incorrect argument!')
-        
+
         return x
 
 
 class ConvBlock5x5(nn.Module):
     def __init__(self, in_channels, out_channels):
-        
+
         super(ConvBlock5x5, self).__init__()
-        
-        self.conv1 = nn.Conv2d(in_channels=in_channels, 
+
+        self.conv1 = nn.Conv2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=(5, 5), stride=(1, 1),
                               padding=(2, 2), bias=False)
-                              
+
         self.bn1 = nn.BatchNorm2d(out_channels)
 
         self.init_weight()
-        
+
     def init_weight(self):
         init_layer(self.conv1)
         init_bn(self.bn1)
 
-        
+
     def forward(self, input, pool_size=(2, 2), pool_type='avg'):
-        
+
         x = input
         x = F.relu_(self.bn1(self.conv1(x)))
         if pool_type == 'max':
@@ -107,27 +108,27 @@ def forward(self, input, pool_size=(2, 2), pool_type='avg'):
             x = x1 + x2
         else:
             raise Exception('Incorrect argument!')
-        
+
         return x
 
 
 class AttBlock(nn.Module):
     def __init__(self, n_in, n_out, activation='linear', temperature=1.):
         super(AttBlock, self).__init__()
-        
+
         self.activation = activation
         self.temperature = temperature
         self.att = nn.Conv1d(in_channels=n_in, out_channels=n_out, kernel_size=1, stride=1, padding=0, bias=True)
         self.cla = nn.Conv1d(in_channels=n_in, out_channels=n_out, kernel_size=1, stride=1, padding=0, bias=True)
-        
+
         self.bn_att = nn.BatchNorm1d(n_out)
         self.init_weights()
-        
+
     def init_weights(self):
         init_layer(self.att)
         init_layer(self.cla)
         init_bn(self.bn_att)
-         
+
     def forward(self, x):
         # x: (n_samples, n_in, n_time)
         norm_att = torch.softmax(torch.clamp(self.att(x), -10, 10), dim=-1)
@@ -143,9 +144,9 @@ def nonlinear_transform(self, x):
 
 
 class Cnn14(nn.Module):
-    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin, 
-        fmax, classes_num):
-        
+    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
+        fmax, classes_num, enable_fusion=False, fusion_type='None'):
+
         super(Cnn14, self).__init__()
 
         window = 'hann'
@@ -155,23 +156,29 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
         amin = 1e-10
         top_db = None
 
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+
         # Spectrogram extractor
-        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size, 
-            win_length=window_size, window=window, center=center, pad_mode=pad_mode, 
+        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size,
+            win_length=window_size, window=window, center=center, pad_mode=pad_mode,
             freeze_parameters=True)
 
         # Logmel feature extractor
-        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size, 
-            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db, 
+        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size,
+            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db,
             freeze_parameters=True)
 
         # Spec augmenter
-        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2, 
+        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2,
             freq_drop_width=8, freq_stripes_num=2)
 
         self.bn0 = nn.BatchNorm2d(64)
 
-        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
+        if (self.enable_fusion) and (self.fusion_type == 'channel_map'):
+            self.conv_block1 = ConvBlock(in_channels=4, out_channels=64)
+        else:
+            self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
         self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
         self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
         self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)
@@ -180,32 +187,116 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
 
         self.fc1 = nn.Linear(2048, 2048, bias=True)
         self.fc_audioset = nn.Linear(2048, classes_num, bias=True)
-        
+
+        if (self.enable_fusion) and (self.fusion_type in ['daf_1d','aff_1d','iaff_1d']):
+            self.mel_conv1d = nn.Sequential(
+                nn.Conv1d(64, 64, kernel_size=5, stride=3, padding=2),
+                nn.BatchNorm1d(64) # No Relu
+            )
+            if self.fusion_type == 'daf_1d':
+                self.fusion_model = DAF()
+            elif self.fusion_type == 'aff_1d':
+                self.fusion_model = AFF(channels=64, type='1D')
+            elif self.fusion_type == 'iaff_1d':
+                self.fusion_model = iAFF(channels=64, type='1D')
+
+        if (self.enable_fusion) and (self.fusion_type in ['daf_2d','aff_2d','iaff_2d']):
+            self.mel_conv2d = nn.Sequential(
+                nn.Conv2d(1, 64, kernel_size=(5,5), stride=(6, 2), padding=(2,2)),
+                nn.BatchNorm2d(64),
+                nn.ReLU(inplace=True)
+            )
+
+            if self.fusion_type == 'daf_2d':
+                self.fusion_model = DAF()
+            elif self.fusion_type == 'aff_2d':
+                self.fusion_model = AFF(channels=64, type='2D')
+            elif self.fusion_type == 'iaff_2d':
+                self.fusion_model = iAFF(channels=64, type='2D')
         self.init_weight()
 
     def init_weight(self):
         init_bn(self.bn0)
         init_layer(self.fc1)
         init_layer(self.fc_audioset)
- 
-    def forward(self, input, mixup_lambda=None):
+
+    def forward(self, input, mixup_lambda=None, device=None):
         """
         Input: (batch_size, data_length)"""
-        x = self.spectrogram_extractor(input)   # (batch_size, 1, time_steps, freq_bins)
-        x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
 
-        x = x.transpose(1, 3)
-        x = self.bn0(x)
-        x = x.transpose(1, 3)
-        
+        if self.enable_fusion and input["longer"].sum() == 0:
+            # if no audio is longer than 10s, then randomly select one audio to be longer
+            input["longer"][torch.randint(0, input["longer"].shape[0], (1,))] = True
+
+        if not self.enable_fusion:
+            x = self.spectrogram_extractor(input['waveform'].to(device=device, non_blocking=True))   # (batch_size, 1, time_steps, freq_bins)
+            x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
+
+            x = x.transpose(1, 3)
+            x = self.bn0(x)
+            x = x.transpose(1, 3)
+        else:
+            longer_list = input["longer"].to(device=device, non_blocking=True)
+            x = input["mel_fusion"].to(device=device, non_blocking=True)
+            longer_list_idx = torch.where(longer_list)[0]
+            x = x.transpose(1, 3)
+            x = self.bn0(x)
+            x = x.transpose(1, 3)
+            if self.fusion_type in ['daf_1d','aff_1d','iaff_1d']:
+                new_x = x[:,0:1,:,:].clone().contiguous()
+                # local processing
+                if len(longer_list_idx) > 0:
+                    fusion_x_local = x[longer_list_idx,1:,:,:].clone().contiguous()
+                    FB,FC,FT,FF = fusion_x_local.size()
+                    fusion_x_local = fusion_x_local.view(FB * FC, FT, FF)
+                    fusion_x_local = torch.permute(fusion_x_local, (0,2,1)).contiguous()
+                    fusion_x_local = self.mel_conv1d(fusion_x_local)
+                    fusion_x_local = fusion_x_local.view(FB,FC,FF,fusion_x_local.size(-1))
+                    fusion_x_local = torch.permute(fusion_x_local, (0,2,1,3)).contiguous().flatten(2)
+                    if fusion_x_local.size(-1) < FT:
+                        fusion_x_local = torch.cat([fusion_x_local, torch.zeros((FB,FF,FT- fusion_x_local.size(-1)), device=device)], dim=-1)
+                    else:
+                        fusion_x_local = fusion_x_local[:,:,:FT]
+                    # 1D fusion
+                    new_x = new_x.squeeze(1).permute((0,2,1)).contiguous()
+                    new_x[longer_list_idx] = self.fusion_model(new_x[longer_list_idx], fusion_x_local)
+                    x = new_x.permute((0,2,1)).contiguous()[:,None,:,:]
+                else:
+                    x = new_x
+            elif self.fusion_type in ['daf_2d','aff_2d','iaff_2d','channel_map']:
+                x = x # no change
+
         if self.training:
             x = self.spec_augmenter(x)
-
         # Mixup on spectrogram
         if self.training and mixup_lambda is not None:
             x = do_mixup(x, mixup_lambda)
+        if (self.enable_fusion) and (self.fusion_type in ['daf_2d','aff_2d','iaff_2d']):
+            global_x = x[:,0:1,:,:]
+
+             # global processing
+            B, C, H, W = global_x.shape
+            global_x = self.conv_block1(global_x, pool_size=(2, 2), pool_type='avg')
+            if len(longer_list_idx) > 0:
+                local_x = x[longer_list_idx,1:,:,:].contiguous()
+                TH = global_x.size(-2)
+                # local processing
+                B, C, H, W = local_x.shape
+                local_x = local_x.view(B*C,1,H,W)
+                local_x = self.mel_conv2d(local_x)
+                local_x = local_x.view(B,C,local_x.size(1),local_x.size(2),local_x.size(3))
+                local_x = local_x.permute((0,2,1,3,4)).contiguous().flatten(2,3)
+                TB,TC,_,TW = local_x.size()
+                if local_x.size(-2) < TH:
+                    local_x = torch.cat([local_x, torch.zeros((TB,TC,TH-local_x.size(-2),TW), device=global_x.device)], dim=-2)
+                else:
+                    local_x = local_x[:,:,:TH,:]
+
+                global_x[longer_list_idx] = self.fusion_model(global_x[longer_list_idx],local_x)
+            x = global_x
+        else:
+            x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
 
-        x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
@@ -218,15 +309,15 @@ def forward(self, input, mixup_lambda=None):
         x = self.conv_block6(x, pool_size=(1, 1), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = torch.mean(x, dim=3)
-        
+
         latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x = latent_x1 + latent_x2
         latent_x = latent_x.transpose(1, 2)
         latent_x = F.relu_(self.fc1(latent_x))
         latent_output = interpolate(latent_x, 32)
-        
-        
+
+
         (x1, _) = torch.max(x, dim=2)
         x2 = torch.mean(x, dim=2)
         x = x1 + x2
@@ -234,15 +325,15 @@ def forward(self, input, mixup_lambda=None):
         x = F.relu_(self.fc1(x))
         embedding = F.dropout(x, p=0.5, training=self.training)
         clipwise_output = torch.sigmoid(self.fc_audioset(x))
-        
+
         output_dict = {'clipwise_output': clipwise_output, 'embedding': embedding, 'fine_grained_embedding': latent_output}
         return output_dict
 
 
 class Cnn6(nn.Module):
-    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin, 
-        fmax, classes_num):
-        
+    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
+        fmax, classes_num, enable_fusion=False, fusion_type='None'):
+
         super(Cnn6, self).__init__()
 
         window = 'hann'
@@ -252,18 +343,21 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
         amin = 1e-10
         top_db = None
 
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+
         # Spectrogram extractor
-        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size, 
-            win_length=window_size, window=window, center=center, pad_mode=pad_mode, 
+        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size,
+            win_length=window_size, window=window, center=center, pad_mode=pad_mode,
             freeze_parameters=True)
 
         # Logmel feature extractor
-        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size, 
-            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db, 
+        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size,
+            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db,
             freeze_parameters=True)
 
         # Spec augmenter
-        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2, 
+        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2,
             freq_drop_width=8, freq_stripes_num=2)
 
         self.bn0 = nn.BatchNorm2d(64)
@@ -275,32 +369,32 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
 
         self.fc1 = nn.Linear(512, 512, bias=True)
         self.fc_audioset = nn.Linear(512, classes_num, bias=True)
-        
+
         self.init_weight()
 
     def init_weight(self):
         init_bn(self.bn0)
         init_layer(self.fc1)
         init_layer(self.fc_audioset)
- 
-    def forward(self, input, mixup_lambda=None):
+
+    def forward(self, input, mixup_lambda=None, device=None):
         """
         Input: (batch_size, data_length)"""
 
         x = self.spectrogram_extractor(input)   # (batch_size, 1, time_steps, freq_bins)
         x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
-        
+
         x = x.transpose(1, 3)
         x = self.bn0(x)
         x = x.transpose(1, 3)
-        
+
         if self.training:
             x = self.spec_augmenter(x)
 
         # Mixup on spectrogram
         if self.training and mixup_lambda is not None:
             x = do_mixup(x, mixup_lambda)
-        
+
         x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')
@@ -310,14 +404,14 @@ def forward(self, input, mixup_lambda=None):
         x = self.conv_block4(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = torch.mean(x, dim=3)
-        
+
         latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x = latent_x1 + latent_x2
         latent_x = latent_x.transpose(1, 2)
         latent_x = F.relu_(self.fc1(latent_x))
         latent_output = interpolate(latent_x, 16)
-        
+
         (x1, _) = torch.max(x, dim=2)
         x2 = torch.mean(x, dim=2)
         x = x1 + x2
@@ -325,16 +419,16 @@ def forward(self, input, mixup_lambda=None):
         x = F.relu_(self.fc1(x))
         embedding = F.dropout(x, p=0.5, training=self.training)
         clipwise_output = torch.sigmoid(self.fc_audioset(x))
-        
+
         output_dict = {'clipwise_output': clipwise_output, 'embedding': embedding, 'fine_grained_embedding': latent_output}
 
         return output_dict
 
 
 class Cnn10(nn.Module):
-    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin, 
-        fmax, classes_num):
-        
+    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
+        fmax, classes_num, enable_fusion=False, fusion_type='None'):
+
         super(Cnn10, self).__init__()
 
         window = 'hann'
@@ -344,18 +438,21 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
         amin = 1e-10
         top_db = None
 
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+
         # Spectrogram extractor
-        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size, 
-            win_length=window_size, window=window, center=center, pad_mode=pad_mode, 
+        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size,
+            win_length=window_size, window=window, center=center, pad_mode=pad_mode,
             freeze_parameters=True)
 
         # Logmel feature extractor
-        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size, 
-            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db, 
+        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size,
+            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db,
             freeze_parameters=True)
 
         # Spec augmenter
-        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2, 
+        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2,
             freq_drop_width=8, freq_stripes_num=2)
 
         self.bn0 = nn.BatchNorm2d(64)
@@ -368,32 +465,32 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
 
         self.fc1 = nn.Linear(1024, 1024, bias=True)
         self.fc_audioset = nn.Linear(1024, classes_num, bias=True)
-        
+
         self.init_weight()
 
     def init_weight(self):
         init_bn(self.bn0)
         init_layer(self.fc1)
         init_layer(self.fc_audioset)
- 
-    def forward(self, input, mixup_lambda=None):
+
+    def forward(self, input, mixup_lambda=None, device=None):
         """
         Input: (batch_size, data_length)"""
 
         x = self.spectrogram_extractor(input)   # (batch_size, 1, time_steps, freq_bins)
         x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
-        
+
         x = x.transpose(1, 3)
         x = self.bn0(x)
         x = x.transpose(1, 3)
-        
+
         if self.training:
             x = self.spec_augmenter(x)
 
         # Mixup on spectrogram
         if self.training and mixup_lambda is not None:
             x = do_mixup(x, mixup_lambda)
-        
+
         x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')
@@ -405,14 +502,14 @@ def forward(self, input, mixup_lambda=None):
         x = self.conv_block5(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = torch.mean(x, dim=3)
-        
+
         latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x = latent_x1 + latent_x2
         latent_x = latent_x.transpose(1, 2)
         latent_x = F.relu_(self.fc1(latent_x))
         latent_output = interpolate(latent_x, 32)
-        
+
         (x1, _) = torch.max(x, dim=2)
         x2 = torch.mean(x, dim=2)
         x = x1 + x2
@@ -420,13 +517,13 @@ def forward(self, input, mixup_lambda=None):
         x = F.relu_(self.fc1(x))
         embedding = F.dropout(x, p=0.5, training=self.training)
         clipwise_output = torch.sigmoid(self.fc_audioset(x))
-        
+
         output_dict = {'clipwise_output': clipwise_output, 'embedding': embedding, 'fine_grained_embedding': latent_output}
 
         return output_dict
 
 
-def create_pann_model(audio_cfg):
+def create_pann_model(audio_cfg, enable_fusion=False, fusion_type='None'):
     try:
         ModelProto = eval(audio_cfg.model_name)
         model = ModelProto(
@@ -436,9 +533,11 @@ def create_pann_model(audio_cfg):
             mel_bins = audio_cfg.mel_bins,
             fmin = audio_cfg.fmin,
             fmax = audio_cfg.fmax,
-            classes_num = audio_cfg.class_num
+            classes_num = audio_cfg.class_num,
+            enable_fusion = enable_fusion,
+            fusion_type = fusion_type
         )
         return model
     except:
         raise RuntimeError(f'Import Model for {audio_cfg.model_name} not found, or the audio cfg parameters are not enough.')
-        
+
diff --git a/src/open_clip/utils.py b/src/open_clip/utils.py
index d7a8fed..ad9801c 100644
--- a/src/open_clip/utils.py
+++ b/src/open_clip/utils.py
@@ -22,9 +22,20 @@
     "wesoundeffects": ["train", "test"],
     "MACS": ["train", "test"],
     "freesound": ["train", "test"],
-    "FSD50K": ["train", "test"],
+    "FSD50K": ["train", "test", "valid"],
+    "fsd50k_class_label": ["train", "test", "valid"],
     "esc50": ["train", "test"],
     "audiostock": ["train", "test"],
+    "freesound_no_overlap_noesc50": ["train", "test"],
+    "epidemic_sound_effects": ["train", "test"],
+    "VGGSound": ["train", "test"],
+    "urbansound8k_class_label": ["train", "test"],
+    "audioset_t5": ["balanced_train", "unbalanced_train", "eval"],
+    "epidemic_sound_effects_t5": ["train", "test"],
+    "WavText5K": ["train", "test"],
+    "esc50_no_overlap": ["train", "test"],
+    "usd8k_no_overlap": ["train", "test"],
+    "fsd50k_200_class_label": ["train", "test", "valid"]
 }
 
 
@@ -85,46 +96,39 @@ def get_tar_path_from_dataset_name(
         islocal,
         dataset_path,
         proportion=1,
+        full_dataset=None
 ):
     """
     Get tar path from dataset name and type
     """
-    if islocal:
-        output = []
-        for n in dataset_names:
-            for s in dataset_types:
-                tmp = []
+    output = []
+    for n in dataset_names:
+        if full_dataset is not None and n in full_dataset:
+            current_dataset_types = dataset_split[n]
+        else:
+            current_dataset_types = dataset_types
+        for s in current_dataset_types:
+            tmp = []
+            if islocal:
                 sizefilepath_ = f"{dataset_path}/{n}/{s}/sizes.json"
                 if not os.path.exists(sizefilepath_):
                     sizefilepath_ = f"./json_files/{n}/{s}/sizes.json"
-                if not os.path.exists(sizefilepath_):
-                    continue
-                sizes = json.load(open(sizefilepath_, "r"))
-                for k in sizes.keys():
-                    tmp.append(f"{dataset_path}/{n}/{s}/{k}")
-                if proportion != 1:
-                    tmp = random.sample(tmp, int(proportion * len(tmp)))
-                output.append(tmp)
-        return sum(output, [])
-    else:
-
-        output = []
-        for n in dataset_names:
-            for s in dataset_types:
-                tmp = []
+            else:
                 sizefilepath_ = f"./json_files/{n}/{s}/sizes.json"
-                if not os.path.exists(sizefilepath_):
-                    continue
-                sizes = json.load(open(sizefilepath_, "r"))
-                for k in sizes.keys():
+            if not os.path.exists(sizefilepath_):
+                continue
+            sizes = json.load(open(sizefilepath_, "r"))
+            for k in sizes.keys():
+                if islocal:
+                    tmp.append(f"{dataset_path}/{n}/{s}/{k}")
+                else:
                     tmp.append(
                         f"pipe:aws s3 --cli-connect-timeout 0 cp s3://s-laion-audio/webdataset_tar/{n}/{s}/{k} -"
                     )
-                    # TODO: add dataset_path to remote dataset in the future.
-                if proportion != 1:
-                    tmp = random.sample(tmp, int(proportion * len(tmp)))
-                output.append(tmp)
-        return sum(output, [])
+            if proportion != 1:
+                tmp = random.sample(tmp, int(proportion * len(tmp)))
+            output.append(tmp)
+    return sum(output, [])
 
 
 def get_tar_path_from_txts(txt_path, islocal, proportion=1):
@@ -343,7 +347,10 @@ def load_class_label(path):
     #     val = Array('i', out.values(), lock=False)
     #     return (key, val)
 
+
 from torch import optim
+
+
 def get_optimizer(params, lr, betas, eps, momentum, optimizer_name):
     if optimizer_name.lower() == "adamw":
         optimizer = optim.AdamW(
@@ -353,6 +360,10 @@ def get_optimizer(params, lr, betas, eps, momentum, optimizer_name):
         optimizer = optim.SGD(
             params, lr=lr, momentum=momentum
         )
+    elif optimizer_name.lower() == "adam":
+        optimizer = optim.Adam(
+            params, lr=lr, betas=betas, eps=eps
+        )
     else:
         raise ValueError("optimizer name is not correct")
     return optimizer
diff --git a/src/training/data.py b/src/training/data.py
index be8063b..a833c4d 100644
--- a/src/training/data.py
+++ b/src/training/data.py
@@ -14,6 +14,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 import torchvision.datasets as datasets
+import torchvision.transforms
 import webdataset as wds
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler
@@ -44,7 +45,7 @@
     from open_clip import tokenize
 
     def tokenizer(text):
-        return tokenize(text)[0, :]
+        return tokenize(text).squeeze(0)
 
 elif args.tmodel == "bert":
     from transformers import BertTokenizer
@@ -97,6 +98,11 @@ def int16_to_float32(x):
     return (x / 32767.0).astype(np.float32)
 
 
+def float32_to_int16(x):
+    x = np.clip(x, a_min=-1., a_max=1.)
+    return (x * 32767.).astype(np.int16)
+
+
 # For Toy Dataset
 class ToyDataset(Dataset):
     def __init__(self, index_path, ipc, config, eval_mode=False):
@@ -218,6 +224,11 @@ def __getitem__(self, index):
         #             target[add_key] = 1.0
 
         # missing the text input
+        mel_spec = get_mel(torch.from_numpy(waveform), self.audio_cfg)[None,:,:]
+        mel_spec = torch.cat([mel_spec, mel_spec.clone(),mel_spec.clone(),mel_spec.clone()], dim=0).cpu().numpy()
+        longer = random.choice([True, False])
+        if longer == False:
+            mel_spec[1:,:,:] = 0.0
         data_dict = {
             "hdf5_path": hdf5_path,
             "index_in_hdf5": r_idx,
@@ -225,6 +236,8 @@ def __getitem__(self, index):
             "waveform": waveform,
             "class_label": target,
             "text": text,
+            "longer": longer,
+            "mel_fusion": mel_spec
         }
         return data_dict
 
@@ -416,20 +429,61 @@ def sample_prop(sizefile, inputs, proportion, is_local=True):
         sampled_size_dict,
     )
 
+def get_mel(audio_data,audio_cfg):
+    # mel shape: (n_mels, T)
+    mel = torchaudio.transforms.MelSpectrogram(
+        sample_rate=audio_cfg['sample_rate'],
+        n_fft=audio_cfg['window_size'],
+        win_length=audio_cfg['window_size'],
+        hop_length=audio_cfg['hop_size'],
+        center=True,
+        pad_mode="reflect",
+        power=2.0,
+        norm=None,
+        onesided=True,
+        n_mels=64,
+        f_min=audio_cfg['fmin'],
+        f_max=audio_cfg['fmax']
+    )(audio_data)
+    # Align to librosa:
+    # librosa_melspec = librosa.feature.melspectrogram(
+    #     waveform,
+    #     sr=audio_cfg['sample_rate'],
+    #     n_fft=audio_cfg['window_size'],
+    #     hop_length=audio_cfg['hop_size'],
+    #     win_length=audio_cfg['window_size'],
+    #     center=True,
+    #     pad_mode="reflect",
+    #     power=2.0,
+    #     n_mels=64,
+    #     norm=None,
+    #     htk=True,
+    #     f_min=audio_cfg['fmin'],
+    #     f_max=audio_cfg['fmax']
+    # )
+    # we use log mel spectrogram as input
+    mel = torchaudio.transforms.AmplitudeToDB(top_db=None)(mel)
+    return mel.T  # (T, n_mels)
+
 
 def preprocess(
     sample,
     audio_ext,
     text_ext,
     max_len,
+    audio_cfg,
     class_index_dict=None,
-    data_filling="repeat"
+    data_filling="pad",
+    data_truncating="rand_trunc",
+    text_augment_selection=None,
 ):
     """
     Preprocess a single sample for wdsdataloader.
     """
     audio_data, orig_sr = sf.read(io.BytesIO(sample[audio_ext]))
+    audio_data = int16_to_float32(float32_to_int16(audio_data))
     audio_data = torch.tensor(audio_data).float()
+
     # TODO: (yusong) to be include in the future
     # # if torchaudio not installed, use soundfile to load audio
     # if torchaudio is None:
@@ -445,37 +499,95 @@ def preprocess(
     #         audio_data, orig_sr = torchaudio.load(fname)
     #         audio_data = audio_data[0, :].float()
 
-    if len(audio_data) > max_len:  # random clip if too long
-        overflow = len(audio_data) - max_len
-        idx = np.random.randint(0, overflow + 1)
-        audio_data = audio_data[idx : idx + max_len]
-    else:  # padding if too short
-        if data_filling == "repeatpad":
-            with torch.no_grad():
-                n_repeat = int(max_len/len(audio_data))
-                audio_data = audio_data.repeat(n_repeat)
-                # audio_data = audio_data.unsqueeze(0).unsqueeze(0).unsqueeze(0)
-                # audio_data = F.interpolate(audio_data,size=max_len,mode="bicubic")[0,0,0]
-                audio_data = F.pad(
-                    audio_data,
-                    (0, max_len - len(audio_data)),
-                    mode="constant",
-                    value=0,
+    with torch.no_grad():
+        if len(audio_data) > max_len:
+            if data_truncating == "rand_trunc":
+                longer = torch.tensor([True])
+            elif data_truncating == "fusion":
+                # fusion
+                mel = get_mel(audio_data, audio_cfg)
+                # split to three parts
+                chunk_frames = max_len // audio_cfg['hop_size']+1  # the +1 related to how the spectrogram is computed
+                total_frames = mel.shape[0]
+                if chunk_frames == total_frames:
+                    # there is a corner case where the audio length is
+                    # larger than max_len but smaller than max_len+hop_size.
+                    # In this case, we just use the whole audio.
+                    mel_fusion = torch.stack([mel, mel, mel, mel], dim=0)
+                    sample["mel_fusion"] = mel_fusion
+                    longer = torch.tensor([False])
+                else:
+                    ranges = np.array_split(list(range(0, total_frames-chunk_frames+1)), 3)
+                    # print('total_frames-chunk_frames:', total_frames-chunk_frames,
+                    #       'len(audio_data):', len(audio_data),
+                    #       'chunk_frames:', chunk_frames,
+                    #       'total_frames:', total_frames)
+                    if len(ranges[1]) == 0:
+                        # if the audio is too short, we just use the first chunk
+                        ranges[1] = [0]
+                    if len(ranges[2]) == 0:
+                        # if the audio is too short, we just use the first chunk
+                        ranges[2] = [0]
+                    # randomly choose index for each part
+                    idx_front = np.random.choice(ranges[0])
+                    idx_middle = np.random.choice(ranges[1])
+                    idx_back = np.random.choice(ranges[2])
+                    # select mel
+                    mel_chunk_front = mel[idx_front:idx_front+chunk_frames, :]
+                    mel_chunk_middle = mel[idx_middle:idx_middle+chunk_frames, :]
+                    mel_chunk_back = mel[idx_back:idx_back+chunk_frames, :]
+
+                    # shrink the mel
+                    mel_shrink = torchvision.transforms.Resize(size=[chunk_frames, 64])(mel[None])[0]
+                    # logging.info(f"mel_shrink.shape: {mel_shrink.shape}")
+
+                    # stack
+                    mel_fusion = torch.stack([mel_chunk_front, mel_chunk_middle, mel_chunk_back, mel_shrink], dim=0)
+                    sample["mel_fusion"] = mel_fusion
+                    longer = torch.tensor([True])
+            else:
+                raise NotImplementedError(
+                    f"data_truncating {data_truncating} not implemented"
                 )
-        elif data_filling == "pad":
-            audio_data = F.pad(
-                audio_data,
-                (0, max_len - len(audio_data)),
-                mode="constant",
-                value=0,
-            )
-        elif data_filling == "repeat":
-            n_repeat = int(max_len/len(audio_data))
-            audio_data = audio_data.repeat(n_repeat+1)[:max_len]
-        else:
-            raise NotImplementedError(
-                f"data_filling {data_filling} not implemented"
-            )
+            # random crop to max_len (for compatibility)
+            overflow = len(audio_data) - max_len
+            idx = np.random.randint(0, overflow + 1)
+            audio_data = audio_data[idx: idx + max_len]
+
+        else:  # padding if too short
+            if len(audio_data) < max_len:  # do nothing if equal
+                if data_filling == "repeatpad":
+                    n_repeat = int(max_len/len(audio_data))
+                    audio_data = audio_data.repeat(n_repeat)
+                    # audio_data = audio_data.unsqueeze(0).unsqueeze(0).unsqueeze(0)
+                    # audio_data = F.interpolate(audio_data,size=max_len,mode="bicubic")[0,0,0]
+                    audio_data = F.pad(
+                        audio_data,
+                        (0, max_len - len(audio_data)),
+                        mode="constant",
+                        value=0,
+                    )
+                elif data_filling == "pad":
+                    audio_data = F.pad(
+                        audio_data,
+                        (0, max_len - len(audio_data)),
+                        mode="constant",
+                        value=0,
+                    )
+                elif data_filling == "repeat":
+                    n_repeat = int(max_len/len(audio_data))
+                    audio_data = audio_data.repeat(n_repeat+1)[:max_len]
+                else:
+                    raise NotImplementedError(
+                        f"data_filling {data_filling} not implemented"
+                    )
+            if data_truncating == 'fusion':
+                mel = get_mel(audio_data, audio_cfg)
+                mel_fusion = torch.stack([mel, mel, mel, mel], dim=0)
+                sample["mel_fusion"] = mel_fusion
+            longer = torch.tensor([False])
+
+    sample["longer"] = longer
 
     sample["waveform"] = audio_data
     del sample[audio_ext]
@@ -484,7 +596,28 @@ def preprocess(
         json_dict_raw = json.loads(sample[text_ext].decode("utf-8"))
     except:
         print("sample[__url__]:", sample["__url__"])
-    texts = json_dict_raw["text"]
+
+    # For selecting augmented text from dataset
+    if text_augment_selection is None or text_augment_selection == "none":
+        texts = json_dict_raw["text"]
+    elif text_augment_selection == "all":
+        if "text_augment_all" in json_dict_raw.keys():
+            texts = json_dict_raw["text_augment_all"]
+        else:
+            texts = json_dict_raw["text"]
+    elif text_augment_selection == "augment_only":
+        if "text_augment_all" in json_dict_raw.keys():
+            if json_dict_raw["text_augment_t5"] is None:
+                texts = json_dict_raw["text"]
+            else:
+                texts = json_dict_raw["text_augment_t5"]
+        else:
+            texts = json_dict_raw["text"]
+    else:
+        raise NotImplementedError(
+            f"text_augment_selection {text_augment_selection} not implemented"
+        )
+    sample["full_text"] = texts
 
     if isinstance(texts, list) and isinstance(texts[0], str) and len(texts) > 1:
         texts = random.choice(texts)
@@ -580,6 +713,7 @@ def get_wds_dataset(
 
     pipeline = [wds.SimpleShardList(input_shards)]
     # at this point we have an iterator over all the shards
+    # TODO: (yusong): add a if statement of distributed. If not, we don't need to split_by_node
     if is_train or args.parallel_eval:
         pipeline.extend(
             [
@@ -615,8 +749,11 @@ def get_wds_dataset(
                 audio_ext=audio_ext,
                 text_ext=text_ext,
                 max_len=max_len,
+                audio_cfg=model_cfg['audio_cfg'],
                 class_index_dict=copy.deepcopy(args.class_index_dict),
                 data_filling=args.data_filling,
+                data_truncating=args.data_truncating,
+                text_augment_selection=args.text_augment_selection,
             )
         ),
     )
@@ -792,13 +929,25 @@ def get_data(args, model_cfg):
             islocal=not args.remotedata,
             proportion=args.dataset_proportion,
             dataset_path=args.datasetpath,
+            full_dataset=args.full_train_dataset,
         )
+
+        if args.full_train_dataset is None:
+            args.full_train_dataset = []
+        if args.exclude_eval_dataset is None:
+            args.exclude_eval_dataset = []
+        excluded_eval_datasets = args.full_train_dataset + args.exclude_eval_dataset
+
+        val_dataset_names = [n for n in args.datasetnames if n not in excluded_eval_datasets] \
+            if excluded_eval_datasets else args.datasetnames
+        args.val_dataset_names = val_dataset_names
         args.val_data = get_tar_path_from_dataset_name(
-            args.datasetnames,
+            val_dataset_names,
             ["valid", "test", "eval"],
             islocal=not args.remotedata,
             proportion=1,
             dataset_path=args.datasetpath,
+            full_dataset=None,
         )
 
     if args.train_data:
diff --git a/src/training/lp_main.py b/src/training/lp_main.py
index e552c09..96c7ed3 100644
--- a/src/training/lp_main.py
+++ b/src/training/lp_main.py
@@ -15,6 +15,7 @@
 import faulthandler
 import pathlib
 import argparse
+import time
 
 try:
     import wandb
@@ -38,7 +39,7 @@
 from training.logger import setup_logging
 from training.scheduler import cosine_lr
 from training.lp_train import train_one_epoch, evaluate
-from open_clip.utils import get_tar_path_from_dataset_name, dataset_split
+from open_clip.utils import get_tar_path_from_dataset_name, dataset_split, get_optimizer
 from open_clip.utils import load_p, load_class_label
 from open_clip.linear_probe import LinearProbe
 
@@ -132,6 +133,12 @@ def random_seed(seed=42, rank=0):
     random.seed(seed + rank)
 
 def config_lp_optimizer(model, data, args):
+    # set wd-related params to 0 if use adam optimizer
+    if args.optimizer == "adam":
+        args.wd = 0
+        args.wd_pretrained = 0
+        args.wd_new = 0
+
     in_clap = (
         lambda n, p: n.startswith("clap_model")
     )
@@ -140,7 +147,7 @@ def config_lp_optimizer(model, data, args):
 
     optimizer = {}
     scheduler = {}
-    
+
     # freeze text encoder
     text_freeze_parameters = [
         p
@@ -150,7 +157,7 @@ def config_lp_optimizer(model, data, args):
         or n.startswith("clap_model.token_embedding")
         or n.startswith("clap_model.ln_final")
     ]
-    
+
     if args.freeze_text:
         logging.info("Freeze Text!!!!")
         for k in text_freeze_parameters:
@@ -166,10 +173,13 @@ def config_lp_optimizer(model, data, args):
         )
         include = lambda n, p: not exclude(n, p)
 
+        # (yusong): we do not split the learning rate anymore
+        # p for n, p in named_parameters if in_clap(n,p) and exclude(n, p) and p.requires_grad
         gain_or_bias_params = [
-            p for n, p in named_parameters if in_clap(n,p) and exclude(n, p) and p.requires_grad
+            p for n, p in named_parameters if exclude(n, p) and p.requires_grad
         ]
-        rest_params = [p for n, p in named_parameters if in_clap(n,p) and include(n, p) and p.requires_grad]
+        # rest_params = [p for n, p in named_parameters if in_clap(n,p) and include(n, p) and p.requires_grad]
+        rest_params = [p for n, p in named_parameters if include(n, p) and p.requires_grad]
 
         if args.train_data is None:
             optimizer = None
@@ -204,7 +214,7 @@ def config_lp_optimizer(model, data, args):
                     if (include(n, p) and p.requires_grad) and (not is_pretrained_params(n))
                 ]
 
-                pretrained_params_optimizer = optim.AdamW(
+                pretrained_params_optimizer = get_optimizer(
                     [
                         {"params": gain_or_bias_pretrained_params, "weight_decay": 0.0},
                         {
@@ -215,6 +225,8 @@ def config_lp_optimizer(model, data, args):
                     lr=args.lr_pretrained,
                     betas=(args.beta1_pretrained, args.beta2_pretrained),
                     eps=args.eps_pretrained,
+                    momentum=args.momentum_pretrained,
+                    optimizer_name=args.optimizer,
                 )
                 pretrained_params_scheduler = cosine_lr(
                     pretrained_params_optimizer,
@@ -223,7 +235,7 @@ def config_lp_optimizer(model, data, args):
                     total_steps,
                 )
 
-                new_params_optimizer = optim.AdamW(
+                new_params_optimizer = get_optimizer(
                     [
                         {"params": gain_or_bias_new_params, "weight_decay": 0.0},
                         {"params": rest_new_params, "weight_decay": args.wd_new},
@@ -231,6 +243,8 @@ def config_lp_optimizer(model, data, args):
                     lr=args.lr_new,
                     betas=(args.beta1_new, args.beta2_new),
                     eps=args.eps_new,
+                    momentum=args.momentum_new,
+                    optimizer_name=args.optimizer,
                 )
                 new_params_scheduler = cosine_lr(
                     new_params_optimizer, args.lr_new, args.warmup, total_steps
@@ -253,8 +267,8 @@ def config_lp_optimizer(model, data, args):
                     hvd.broadcast_optimizer_state(pretrained_params_optimizer, root_rank=0)
                     hvd.broadcast_optimizer_state(new_params_optimizer, root_rank=0)
             else:
-                                
-                optimizer["clap"] = optim.AdamW(
+
+                optimizer["clap"] = get_optimizer(
                     [
                         {"params": gain_or_bias_params, "weight_decay": 0.0},
                         {"params": rest_params, "weight_decay": args.wd},
@@ -262,7 +276,9 @@ def config_lp_optimizer(model, data, args):
                     lr=args.lr,
                     betas=(args.beta1, args.beta2),
                     eps=args.eps,
-                )
+                    momentum=args.momentum,
+                    optimizer_name=args.optimizer,
+            )
                 scheduler["clap"] = cosine_lr(optimizer["clap"], args.lr, args.warmup, total_steps)
 
                 if args.horovod:
@@ -273,15 +289,20 @@ def config_lp_optimizer(model, data, args):
                     hvd.broadcast_optimizer_state(optimizer["clap"], root_rank=0)
 
     # linear probe optimizer
-    lp_params = [p for n,p in named_parameters if (not in_clap(n,p)) and p.requires_grad]
-    lp_optim = optim.AdamW(lp_params, lr = args.lp_lr)
-    optimizer["lp"] = lp_optim
+    else:
+        lp_params = [p for n, p in named_parameters if (not in_clap(n, p)) and p.requires_grad]
+        lp_optim = get_optimizer(lp_params, lr=args.lp_lr, betas=(args.beta1, args.beta2), eps=args.eps, momentum=0.9,
+                                 optimizer_name=args.optimizer)
+        optimizer["lp"] = lp_optim
 
     return optimizer, scheduler, text_freeze_parameters
 
 
 def main():
     args = parse_args()
+
+    time.sleep(args.sleep)
+
     # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
     args.amodel = args.amodel.replace("/", "-")
     # download sizes.json file
@@ -330,11 +351,19 @@ def main():
         os.makedirs(log_base_path, exist_ok=True)
         log_filename = f"out-{args.rank}" if args.log_local else "out.log"
         args.log_path = os.path.join(log_base_path, log_filename)
-        if os.path.exists(args.log_path):
-            print(
-                "Error. Experiment already exists. Use --name {} to specify a new experiment."
-            )
-            return -1
+
+        # avoid log dir in same name:
+        postfix = 0
+        while os.path.exists(args.log_path):
+            postfix += 1
+            log_base_path_new = log_base_path+'-'+str(postfix)
+            os.makedirs(log_base_path_new, exist_ok=True)
+            log_filename = f"out-{args.rank}" if args.log_local else "out.log"
+            args.log_path = os.path.join(log_base_path_new, log_filename)
+            # print(
+            #     "Error. Experiment already exists. Use --name {} to specify a new experiment."
+            # )
+            # return -1
 
     # Set logger
     args.log_level = logging.DEBUG if args.debug else logging.INFO
@@ -396,9 +425,11 @@ def main():
         openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
         skip_params=False,
         pretrained_audio=args.pretrained_audio,
-        pretrained_text=args.pretrained_text
+        pretrained_text=args.pretrained_text,
+        enable_fusion=args.enable_fusion,
+        fusion_type=args.fusion_type
     )
-    
+
     args.lp_out_ch = len(list(args.class_index_dict.keys()))
     # Linear Probe 
     logging.info(f"linear probe using mlp: {args.lp_mlp}")
@@ -410,13 +441,13 @@ def main():
     logging.info(f"linear probe lp_metrics: {args.lp_metrics}")
 
     model = LinearProbe(
-        clap_model, 
-        mlp=args.lp_mlp, freeze=args.lp_freeze, 
+        clap_model,
+        mlp=args.lp_mlp, freeze=args.lp_freeze,
         in_ch=512, out_ch=args.lp_out_ch,
         act=args.lp_act
     ) # in_ch is fixed (i.e., 512)
     model = model.to(device)
-    
+
     if args.horovod:
         with torch.no_grad():
             for param in model.parameters():
@@ -452,7 +483,7 @@ def main():
     if args.trace:
         assert "train" not in data, "Cannot train with traced model"
 
-    
+
     optimizer, scheduler, text_freeze_parameters = config_lp_optimizer(model, data, args)
 
 
diff --git a/src/training/lp_train.py b/src/training/lp_train.py
index a2b32ef..f5a977f 100644
--- a/src/training/lp_train.py
+++ b/src/training/lp_train.py
@@ -15,6 +15,7 @@
     wandb = None
 
 from open_clip import LPLoss, LPMetrics, lp_gather_features
+from open_clip.utils import do_mixup, get_mix_lambda
 from .distributed import is_master
 from .zero_shot import zero_shot_eval
 
@@ -77,12 +78,18 @@ def train_one_epoch(
         else:
             scheduler(step)
 
-        audio = batch['waveform']
+        audio = batch # contains mel_spec, wavform, and longer list
         class_label = batch['class_label']
-
-        audio = audio.to(device=device, non_blocking=True)
+        # audio = audio.to(device=device, non_blocking=True)
         class_label = class_label.to(device=device, non_blocking=True)
 
+        if args.mixup:
+            # https://github.com/RetroCirce/HTS-Audio-Transformer/blob/main/utils.py#L146
+            mix_lambda = torch.from_numpy(get_mix_lambda(0.5, len(audio["waveform"]))).to(device)
+            class_label = do_mixup(class_label, mix_lambda)
+        else:
+            mix_lambda = None
+
         data_time_m.update(time.time() - end)
         if isinstance(optimizer, dict):
             for o_ in optimizer.values():
@@ -91,7 +98,7 @@ def train_one_epoch(
             optimizer.zero_grad()
 
         with autocast():
-            pred = model(audio)
+            pred = model(audio, mix_lambda=mix_lambda, device=device)
             total_loss = loss(pred, class_label)
 
         if isinstance(optimizer, dict):
@@ -135,7 +142,10 @@ def train_one_epoch(
         batch_count = i + 1
 
         if is_master(args) and (i % 100 == 0 or batch_count == num_batches_per_epoch):
-            batch_size = len(audio)
+            if isinstance(audio, dict):
+                batch_size = len(audio["waveform"])
+            else:
+                batch_size = len(audio)
             num_samples = batch_count * batch_size * args.world_size
             samples_per_epoch = dataloader.num_samples
             percent_complete = 100.0 * batch_count / num_batches_per_epoch
@@ -148,7 +158,7 @@ def train_one_epoch(
                     f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
                     f"Data (t): {data_time_m.avg:.3f} "
                     f"Batch (t): {batch_time_m.avg:.3f} "
-                    f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]}" 
+                    f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]}"
                 )
                 log_data = {
                     "loss": loss_m.val,
@@ -185,7 +195,6 @@ def train_one_epoch(
             data_time_m.reset()
     # end for
 
-
 def evaluate(model, data, epoch, args, tb_writer=None, extra_suffix=""):
     metrics = {}
     if not args.parallel_eval:
@@ -223,14 +232,14 @@ def evaluate(model, data, epoch, args, tb_writer=None, extra_suffix=""):
         }
         with torch.no_grad():
             for i, batch in enumerate(dataloader):
-                audio = batch['waveform']
+                audio = batch # contains mel_spec, wavform, and longer list
                 class_label = batch['class_label']
                            
-                audio = audio.to(device=device, non_blocking=True)
+                # audio = audio.to(device=device, non_blocking=True)
                 class_label = class_label.to(device=device, non_blocking=True)
 
                 with autocast():
-                    pred = model(audio)
+                    pred = model(audio, device=device)
                     if args.parallel_eval:
                         pred, class_label = lp_gather_features(pred, class_label, args.world_size, args.horovod)
                     eval_info['pred'].append(pred)
@@ -281,5 +290,3 @@ def evaluate(model, data, epoch, args, tb_writer=None, extra_suffix=""):
         return metrics
     else:
         return metrics
-
-
diff --git a/src/training/main.py b/src/training/main.py
index d937a85..fefc13b 100644
--- a/src/training/main.py
+++ b/src/training/main.py
@@ -244,7 +244,9 @@ def main():
         openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
         skip_params=True,
         pretrained_audio=args.pretrained_audio,
-        pretrained_text=args.pretrained_text
+        pretrained_text=args.pretrained_text,
+        enable_fusion=args.enable_fusion,
+        fusion_type=args.fusion_type
     )
 
     if args.horovod:
@@ -297,10 +299,7 @@ def main():
     text_freeze_parameters = [
         p
         for n, p in named_parameters
-        if n.startswith("text_branch")
-        or n in ["positional_embedding", "text_projection"]
-        or n.startswith("token_embedding")
-        or n.startswith("ln_final")
+        if 'text_branch' in n
     ]
 
     if args.freeze_text:
@@ -313,6 +312,12 @@ def main():
     ]
     rest_params = [p for n, p in named_parameters if include(n, p) and p.requires_grad]
 
+    # set wd-related params to 0 if use adam optimizer
+    if args.optimizer == "adam":
+        args.wd = 0
+        args.wd_pretrained = 0
+        args.wd_new = 0
+
     if args.train_data is None:
         optimizer = None
         scheduler = None
@@ -356,7 +361,7 @@ def main():
                 lr=args.lr_pretrained,
                 betas=(args.beta1_pretrained, args.beta2_pretrained),
                 eps=args.eps_pretrained,
-                momentum=args.momentum_pretraatariined,
+                momentum=args.momentum_pretrained,
                 optimizer_name=args.optimizer,
                 )
             pretrained_params_scheduler = cosine_lr(
diff --git a/src/training/params.py b/src/training/params.py
index 84e7d6e..a527122 100644
--- a/src/training/params.py
+++ b/src/training/params.py
@@ -102,6 +102,18 @@ def parse_args():
         default=None,
         help="If loading webdataset, spedify the dataset names to load. Can be some of these: Clotho, audioset, audiocaps, BBCSoundEffects",
     )
+    parser.add_argument(
+        "--full-train-dataset",
+        nargs="+",
+        default=None,
+        help="Which dataset will be trained with all the subsets. (train+test)",
+    )
+    parser.add_argument(
+        "--exclude-eval-dataset",
+        nargs="+",
+        default=None,
+        help="Which dataset will be excluded with evaluation",
+    )
     parser.add_argument(
         "--datasetinfos",
         nargs="+",
@@ -188,7 +200,7 @@ def parse_args():
         "--wd-pretrained", type=float, default=0.2, help="Weight decay for text."
     )
     parser.add_argument(
-        "--momentum-pretrained", type=float, default=0.2, help="Momentum for text."
+        "--momentum-pretrained", type=float, default=0.9, help="Momentum for text."
     )
     parser.add_argument(
         "--lr-new", type=float, default=None, help="Learning rate for audio."
@@ -206,7 +218,7 @@ def parse_args():
         "--wd-new", type=float, default=0.2, help="Weight decay for audio."
     )
     parser.add_argument(
-        "--momentum-new", type=float, default=0.2, help="Momentum for audio."
+        "--momentum-new", type=float, default=0.9, help="Momentum for audio."
     )
     parser.add_argument(
         "--warmup", type=int, default=10000, help="Number of steps to warmup for."
@@ -473,14 +485,23 @@ def parse_args():
         "--lp-lr", type=float, default=1e-4, help="learning rate of linear probe"
     )
     parser.add_argument(
-        "--kappa", type=float, default=0, help="the kappa in the weighted contrastive loss, default is to turn off the weighted contrastive loss"
+        "--kappa", type=float, default=0,
+        help="the kappa in the weighted contrastive loss, default is to turn off the weighted contrastive loss"
     )
 
     parser.add_argument(
         "--data-filling",
         type=str,
-        default="repeat",
-        help="type of data filling, can be one of the following: repeat, repeatpad, pad",
+        default="pad",
+        help="type of data filling when the audio length is shorter than the max length."
+             "Can be one of the following: repeat, repeatpad, pad",
+    )
+    parser.add_argument(
+        "--data-truncating",
+        type=str,
+        default="rand_trunc",
+        help="type of data truncation when the audio length is longer than the max length."
+             "Can be one of the following: rand_trunc, fusion",
     )
 
     parser.add_argument(
@@ -490,6 +511,45 @@ def parse_args():
         help="Using MLP loss for CLAP model or not",
     )
 
+    parser.add_argument(
+        "--wandb-id",
+        type=str,
+        default=None,
+        help="the id of wandb experiment to restore.",
+    )
+
+    parser.add_argument(
+        "--sleep", type=float, default=0, help="sleep n seconds before start training"
+    )
+
+    # variable length processing
+    parser.add_argument(
+        "--enable-fusion",
+        default=False,
+        action="store_true",
+        help="Enable feature funsion for variable-length data",
+    )
+
+    parser.add_argument(
+        "--fusion-type",
+        type=str,
+        default='None',
+        help="Type is among ['channel_map', 'daf_1d','aff_1d','iaff_1d','daf_2d','aff_2d','iaff_2d']",
+    )
+
+    parser.add_argument(
+        "--mixup",
+        default=False,
+        action="store_true",
+        help="Enable mixup in finetuning training.",
+    )
+    parser.add_argument(
+        "--text-augment-selection",
+        type=str,
+        default=None,
+        help="For selecting levels of augmented text. Type is among ['all', 'augment_only', 'none']",
+    )
+
     args = parser.parse_args()
 
     # If some params are not passed, we use the default values based on model name.
diff --git a/src/training/train.py b/src/training/train.py
index bf07e57..f3e41b1 100644
--- a/src/training/train.py
+++ b/src/training/train.py
@@ -79,16 +79,15 @@ def train_one_epoch(
 
     for i, batch in enumerate(dataloader):
         # logging.info(f"batch {i} of {num_batches_per_epoch}")
-        # logging.info(model.state_dict()['module.text_branch.resblocks.4.attn.in_proj_bias'][:10]) # For debug
         step = num_batches_per_epoch * epoch + i
         if isinstance(scheduler, dict):
             for s in scheduler.values():
                 s(step)
         else:
             scheduler(step)
-        audios = batch['waveform']
+        audios = batch  # contains mel_spec, wavform, and longer list
         texts = batch['text']
-        audios = audios.to(device=device, non_blocking=True)
+        # audios = audios.to(device=device, non_blocking=True)
         # texts = texts.to(device=device, non_blocking=True)
 
         data_time_m.update(time.time() - end)
@@ -164,7 +163,10 @@ def train_one_epoch(
         end = time.time()
         batch_count = i + 1
         if is_master(args) and (i % 100 == 0 or batch_count == num_batches_per_epoch):
-            batch_size = len(audios)
+            if isinstance(audios, dict):
+                batch_size = len(audios["waveform"])
+            else:
+                batch_size = len(audios)
             num_samples = batch_count * batch_size * args.world_size
             samples_per_epoch = dataloader.num_samples
             percent_complete = 100.0 * batch_count / num_batches_per_epoch
@@ -276,7 +278,19 @@ def evaluate(model, data, epoch, args, tb_writer=None):
     if is_master(args):
         print('Evaluating...')
     autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
-    if "val" in data and (
+    if args.val_dataset_names == ['Clotho', 'audiocaps']:
+        # if only clotho and audiocaps are used, then we will use a different evaluation function.
+        # This is because in the Clotho and audiocaps valid and test set, there are 5 text for 1 audio.
+        if args.parallel_eval:
+            # (yusong): just a hack here. Don't use parallel eval when evaluating only clotho and audiocaps.
+            raise NotImplementedError("Parallel evaluation not supported for eval only Clotho and audiocaps.")
+        val_metrics_per_dataset = evaluate_clotho_audiocaps(model, data, epoch, args, autocast, device, tb_writer)
+        for m in val_metrics_per_dataset.values():
+            metrics.update(m)
+        if "epoch" not in metrics.keys():
+            metrics.update({"epoch": epoch})
+        metrics = select_top_metric_clotho_audiocaps(metrics, val_metrics_per_dataset, args)
+    elif "val" in data and (
             args.val_frequency
             and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)
     ):
@@ -306,10 +320,10 @@ def evaluate(model, data, epoch, args, tb_writer=None):
         # all_audio_features, all_text_features, all_audio_features_mlp, all_text_features_mlp = [], [], [], []
         with torch.no_grad():
             for i, batch in enumerate(dataloader):
-                audios = batch['waveform']
+                audios = batch  # contains mel_spec, wavform, and longer list
                 texts = batch['text']
-                audios = audios.to(device=device, non_blocking=True)
-                
+                # audios = audios.to(device=device, non_blocking=True)
+
                 all_names = list(set(["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]))
                 for name in all_names:
                     if name not in eval_info.keys():
@@ -428,7 +442,7 @@ def evaluate(model, data, epoch, args, tb_writer=None):
                         f"Eval Epoch: {epoch} [{num_samples} / {samples_per_val}]"
                     )
             if is_master(args):
-                val_metrics_s = {}
+                val_metrics_per_dataset = {}
                 for n in eval_info.keys():
                     if args.clap_mlploss:
                         metrics_single_dataset = get_metrics(
@@ -449,10 +463,10 @@ def evaluate(model, data, epoch, args, tb_writer=None):
                             logit_scale_a=logit_scale_a.cpu(),
                             mlp_loss=args.clap_mlploss
                         )
-                    val_metrics_s[n] = {
+                    val_metrics_per_dataset[n] = {
                         n + "/" + k: v for k, v in metrics_single_dataset.items()
                     }
-                    metrics.update(val_metrics_s[n])
+                    metrics.update(val_metrics_per_dataset[n])
                     if "epoch" not in metrics.keys():
                         metrics.update({"epoch": epoch})
     if is_master(args):
@@ -464,7 +478,7 @@ def evaluate(model, data, epoch, args, tb_writer=None):
             + "\n".join(
                 [
                     "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in m.items()])
-                    for m in val_metrics_s.values()
+                    for m in val_metrics_per_dataset.values()
                 ]
             )
         )
@@ -512,11 +526,11 @@ def get_metrics(
         labels = torch.arange(audio_features.shape[0]).long()
         # Change the loss from two terms into four terms with 2x2 combined CE loss
         total_loss = (
-            F.cross_entropy(a_logits_per_audio, labels)
-            + F.cross_entropy(a_logits_per_text, labels)
-            + F.cross_entropy(t_logits_per_audio, labels)
-            + F.cross_entropy(t_logits_per_text, labels)
-        ) / 4
+                             F.cross_entropy(a_logits_per_audio, labels)
+                             + F.cross_entropy(a_logits_per_text, labels)
+                             + F.cross_entropy(t_logits_per_audio, labels)
+                             + F.cross_entropy(t_logits_per_text, labels)
+                     ) / 4
 
         metrics[f"cumulative_loss"] = total_loss.item()
         metrics[f"num_samples"] = audio_features.shape[0]
@@ -536,23 +550,20 @@ def get_metrics(
         labels = torch.arange(audio_features.shape[0]).long()
         # Change the loss from two terms into four terms with 2x2 combined CE loss
         total_loss = (
-            F.cross_entropy(logits_per_audio, labels)
-            + F.cross_entropy(logits_per_text, labels)
-        ) / 2
+                             F.cross_entropy(logits_per_audio, labels)
+                             + F.cross_entropy(logits_per_text, labels)
+                     ) / 2
 
         metrics[f"cumulative_loss"] = total_loss.item()
         metrics[f"num_samples"] = audio_features.shape[0]
 
-
         logits = {"audio_to_text": logits_per_audio, "text_to_audio": logits_per_text}
-        
+
         ground_truth = torch.arange(len(text_features)).view(-1, 1)
 
     for name, logit in logits.items():
         ranking = torch.argsort(logit, descending=True)
-        preds = torch.where(ranking == ground_truth)[
-            1
-        ]  # (yusong) this line is slow because it uses single thread
+        preds = torch.where(ranking == ground_truth)[1]  # (yusong) this line is slow because it uses single thread
         preds = preds.detach().cpu().numpy()
         metrics[f"{name}_mean_rank"] = preds.mean() + 1
         metrics[f"{name}_median_rank"] = np.floor(np.median(preds)) + 1
@@ -561,6 +572,209 @@ def get_metrics(
         # map@10
         metrics[f"{name}_mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
 
+    return metrics
+
 
+def evaluate_clotho_audiocaps(
+        model, data, epoch, args, autocast, device, tb_writer=None
+):
+    """
+    Adapted from https://github.com/XinhaoMei/audio-text_retrieval/blob/main/tools/utils.py.
+    1. for text-to-audio retrieval, do 5 times and average the results
+    2. for R@1, R@5, R@10 in audio-to-text retrieval, take the best rank among 5 text
+    3. for map@10 in audio-to-text retrieval:
+        3.1: sort the rank of 5 text
+        3.2: exclude the rank >=10 (0-index)
+        3.3: compute the map regarding the remaining ranks: np.mean(np.arange(1, len(ranks)+1) / ranks).
+        (3.3) That is, take the top ranks of 5 text that is < 10, and assign the descending number as ground truth.
+        (3.3) E.g.: the ground truth of first rank of the 5 text should be 1, the second rank should be 2, etc.
+    """
+    # TODO: (yusong) only support single GPU evaluation and only support non-mlp case for now.
+    dataloader = data["val"].dataloader
+    with torch.no_grad():
+        eval_info = {}
+        for i, batch in enumerate(dataloader):
+            audios = batch  # contains mel_spec, wavform, and longer list
+
+            # each item in the list has 5 texts
+            if args.tmodel == "transformer":
+                from open_clip import tokenize
+                texts = [tokenize(t) for t in batch['full_text']]
+                texts = torch.cat(texts)
+            else:
+                from .data import tokenizer
+                texts = [tokenizer(t) for t in batch['full_text']]  # 5 texts for each audio
+                texts = {k: torch.cat([t[k] for t in texts]) for k in texts[0].keys()}  # 5 x batch
+
+            # audios = audios.to(device=device, non_blocking=True)
+
+            all_names = list(set(["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]))
+            for name in all_names:
+                if name not in eval_info.keys():
+                    # we will not use mlp outputs even if args.clap_mlploss=True
+                    eval_info[name] = {
+                        "cumulative_loss": 0.0,
+                        "num_samples": 0,
+                        "all_audio_features": [],
+                        "all_text_features": []
+                    }
+            with autocast():
+                audio_features = model(audios, None, device)
+                text_features = model(None, texts, device)
+                audio_features = F.normalize(audio_features, dim=-1)
+                text_features = F.normalize(text_features, dim=-1)
+
+                all_names = list(set(["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]))
+                for n in all_names:
+                    idx = np.where(
+                        np.array(
+                            ["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]
+                        )
+                        == n
+                    )[0]
+                    eval_info[n]["all_audio_features"].append(
+                        audio_features.cpu().index_select(
+                            0, torch.tensor(idx).long()
+                        )
+                    )
+                    # (yusong) please double-check. This is for selecting 5 text features at once.
+                    # because idx is a list of indices in size of num_samples,
+                    # and text_features is a tensor of size (5*num_samples, dim)
+                    # so we need to select 5 consecutive indices at once for a single index in idx.
+                    eval_info[n]["all_text_features"].append(
+                        text_features.cpu().reshape([-1, 5, text_features.shape[1]]).index_select(
+                            0, torch.tensor(idx).long()
+                        ).reshape([-1, text_features.shape[1]])
+                    )
+
+        val_metrics_all = {}
+
+        for n in eval_info.keys():
+            logit_scale_a, logit_scale_t = model(None, None, device)
+            logit_scale_a = logit_scale_a.cpu()
+
+            audio_features = torch.cat(eval_info[n]["all_audio_features"], dim=0)
+            text_features = torch.cat(eval_info[n]["all_text_features"], dim=0)
+
+            logits_per_audio = (logit_scale_a * audio_features @ text_features.t()).detach().cpu()
+            logits_per_text = logits_per_audio.t().detach().cpu()
+
+            # logits_per_audio shape: [num_samples, num_samples*5]
+            # logits_per_text shape: [num_samples*5, num_samples]
+
+            logging.info(f"dataset {n}, logits_per_audio shape: {logits_per_audio.shape}, "
+                         f"logits_per_text shape: {logits_per_text.shape}")
+
+            metrics = {}
+            num_samples = audio_features.shape[0]
+            metrics[f"num_samples"] = num_samples
+
+            # (yusong) the following code is very important, please double-check:
+            # logits_per_audio.reshape(num_samples, num_samples, 5)[:, :, d]
+            # logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :]
+            # Those two are retrieving one of the 5 text for each audio.
+            labels = torch.arange(audio_features.shape[0]).long()
+            audio_to_text_loss = [
+                F.cross_entropy(
+                    logits_per_audio.reshape(num_samples, num_samples, 5)[:, :, d], labels) for d in range(5)
+            ]
+            text_to_audio_loss = [
+                F.cross_entropy(
+                    logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :], labels) for d in range(5)
+            ]
+            total_loss = (
+                                 np.mean(audio_to_text_loss) + np.mean(text_to_audio_loss)
+                         ) / 2
+
+            metrics[f"cumulative_loss"] = total_loss.item()
+
+            # text to audio: do 5 times
+            pred_text = []
+            for d in range(5):
+                logit = logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :]
+                ground_truth = torch.arange(len(logit)).view(-1, 1)
+                ranking = torch.argsort(logit, descending=True)  # [num_samples, num_samples]
+                preds = torch.where(ranking == ground_truth)[1]
+                pred_text.append(preds.detach().cpu().numpy())
+            pred_text_concat = np.concatenate(pred_text, axis=0)  # [5*num_samples]
+            metrics[f"text_to_audio_mean_rank"] = pred_text_concat.mean() + 1
+            metrics[f"text_to_audio_median_rank"] = np.floor(np.median(pred_text_concat)) + 1
+            for k in [1, 5, 10]:
+                metrics[f"text_to_audio_R@{k}"] = np.mean(pred_text_concat < k)
+            # map@10
+            metrics[f"text_to_audio_mAP@10"] = np.mean(np.where(pred_text_concat < 10, 1 / (pred_text_concat + 1), 0.0))
+
+            # audio to text: take the best result
+            # for audio to text map 10, sort and assign descending ground truth.
+            # see https://github.com/XinhaoMei/audio-text_retrieval/blob/main/tools/utils.py#L103
+            # map@10
+            map_all = []
+            pred_audio_all = []
+            for d in range(num_samples):
+                # logits_per_audio: [num_samples, num_samples*5]
+                logit_single = logits_per_audio[d, :]  # [5*num_samples]
+                # Ground-truth index: [d*5, d*5+1, d*5+2, d*5+3, d*5+4]
+                ranking = torch.argsort(logit_single, descending=True)  # [5*num_samples]
+                # ranking: the index of first match, second match, ...
+                ground_truth = torch.arange(d * 5, d * 5 + 5)[None]
+                all_pred = torch.where(torch.stack([ranking] * 5) == ground_truth.view(-1, 1))[1]
+                min_pred = torch.min(all_pred)
+                pred_audio_all.append(min_pred.detach().cpu().numpy())
+                all_pred_filter = all_pred[all_pred < 10].detach().cpu().numpy()
+                # /5 because we have 5 text, so it means for the text rank >=10 we count as 0.
+                map_single = np.sum((np.arange(1, len(all_pred_filter) + 1) / (all_pred_filter + 1))) / 5
+                map_all.append(map_single)
+            metrics[f"audio_to_text_mAP@10"] = np.mean(map_all)
+            for k in [1, 5, 10]:
+                metrics[f"audio_to_text_R@{k}"] = np.mean(np.array(pred_audio_all) < k)
+
+            val_metrics_all[n] = {
+                n + "/" + k: v for k, v in metrics.items()
+            }
+    return val_metrics_all
+
+
+def calculate_selection_performance_clotho_audiocaps(val_metrics_per_dataset):
+    """
+    Calculate performance for Clotho+AudioCaps for model selection.
+    """
+    selection_performance_all = []
+    for n in val_metrics_per_dataset.keys():
+        selection_performance = (val_metrics_per_dataset[n][f"{n}/audio_to_text_mAP@10"] +
+                                 val_metrics_per_dataset[n][f"{n}/text_to_audio_mAP@10"]) / 2
+        selection_performance_all.append(selection_performance)
+    return np.mean(selection_performance_all)
+
+
+def select_top_metric_clotho_audiocaps(metrics, val_metrics_per_dataset, args):
+    # val_metrics_per_dataset: dict, key: dataset name, value: dict, key: metric name, value: metric value
+    # metrics: dict, key: metric name, value: metric value
+    # Hack: use args to save the top performance
+    if not hasattr(args, "top_selection_performance"):
+        selection_performance = calculate_selection_performance_clotho_audiocaps(val_metrics_per_dataset)
+        # TODO: write the if and else together
+        metric_update = {}
+        for n in val_metrics_per_dataset.keys():
+            for k in val_metrics_per_dataset[n].keys():
+                metric_update[k.split('/')[0] + '-top' + '/' + k.split('/')[1]] = val_metrics_per_dataset[n][k]
+        metric_update['top_selection_performance'] = selection_performance
+        metric_update['top-selection-epoch'] = metrics['epoch']
+        metrics.update(metric_update)
+        args.top_metric = metric_update
+        args.top_selection_performance = selection_performance
+    else:
+        selection_performance_new = calculate_selection_performance_clotho_audiocaps(val_metrics_per_dataset)
+        selection_performance_old = args.top_selection_performance
+        if selection_performance_new > selection_performance_old:
+            metric_update = {}
+            for n in val_metrics_per_dataset.keys():
+                for k in val_metrics_per_dataset[n].keys():
+                    metric_update[k.split('/')[0] + '-top' + '/' + k.split('/')[1]] = val_metrics_per_dataset[n][k]
+            metric_update['top_selection_performance'] = selection_performance_new
+            metric_update['top-selection-epoch'] = metrics['epoch']
+            metrics.update(metric_update)
+            args.top_metric = metric_update
+            args.top_selection_performance = selection_performance_new
+        else:
+            metrics.update(args.top_metric)
     return metrics
-    
\ No newline at end of file


Test patch:
diff --git a/src/training/test_infer.py b/src/training/test_infer.py
index 856456a..2c6b7cc 100644
--- a/src/training/test_infer.py
+++ b/src/training/test_infer.py
@@ -249,7 +249,7 @@ def main():
     ]
 
     # construct data
-    output_dict = model.audio_infer(torch.rand(960000).cuda(), hopsize=120000, key ="fine_grained_embedding")
+    output_dict = model.audio_infer(torch.rand(960000).cuda(), hopsize=120000, key ="fine_grained_embedding", device=device)
     print(output_dict["fine_grained_embedding"].size())
 
     # inference