python pytorchで書かれた関数のリファクタリング

【概要】python pytorch で書かれたある関数(約70行)のコードの改善を手伝ってほしい 【目的】実験に使用します。 【依頼背景】 私が書いたものではないのですが、処理内容を変えずに、見やすいコードに急ぎで書き換える必要があります。for文が多く、ボトルネックになっているようです…。早めに対応頂ける方が望ましいです。 【提案の際のお願い】 数千行もあるうちの1つの関数部分のコードのみを渡すので、それだけで実際に実行することはできません。 どう書けばより美しくなるのか、煮詰まってしまったのでアイデアが欲しいです。 python らしい書き方を教えて頂ければ嬉しいです。