support NTU

2021-09-03 11:06:26 +08:00
parent 0702081dbe
commit 16387a7afe
8 changed files with 236 additions and 23 deletions
--- a/fit/tools/main.py
+++ b/fit/tools/main.py
@ -114,7 +114,8 @@ if __name__ == "__main__":
            meters.reset_early_stop()
            logger.info("avg_loss:{:.4f}".format(meters.avg))

-            # save_pic(res,smpl_layer,file,logger,args.dataset_name,target)
            save_params(res, file, logger, args.dataset_name)
+            save_pic(res,smpl_layer,file,logger,args.dataset_name,target)
+            
            torch.cuda.empty_cache()
    logger.info("Fitting finished! Average loss:     {:.9f}".format(meters.avg))