training code

[MandyMo]

When will the training code be released, can not wait to have a try !!

[rishabh-akridata]

@MandyMo Did you manage to load and infer the model? It's crashing for me because it is consuming too much memory, both cuda and DRAM.

Thanks.