class_docstr: Stores inference context and reuses it in sequence embedding output_dist or result return.
  
  Attributes:
      features KJTList: stores the shards of KJT after input dist.
      features_before_input_dist KJT: stores the original input KJT (before input dist).
      unbucketize_permute_tensor Optional[torch.Tensor]: stores unbucketize tensor, only for RowWise sharding.
class_docstr_tok: ["Stores","inference","context","and","reuses","it","in","sequence","embedding","output_dist","or","result","return",".","Attributes",":","features","KJTList",":","stores","the","shards","of","KJT","after","input","dist",".","features_before_input_dist","KJT",":","stores","the","original","input","KJT","(","before","input","dist",")",".","unbucketize_permute_tensor","Optional","[","torch.Tensor","]",":","stores","unbucketize","tensor",",","only","for","RowWise","sharding","."]
code: def record_stream(self, stream: torch.cuda.streams.Stream) -> None:
      for feature in self.features:
          feature.record_stream(stream)
      if self.features_before_input_dist is not None:
          self.features_before_input_dist.record_stream(stream)
      if self.unbucketize_permute_tensor is not None:
          self.unbucketize_permute_tensor.record_stream(stream)
  
file_path: torchrec/torchrec/distributed/sharding/sequence_sharding.py
imports: from dataclasses import dataclass, field
  from typing import List, Optional
  import torch
  from torchrec.distributed.embedding_sharding import EmbeddingShardingContext
  from torchrec.distributed.embedding_types import KJTList
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
  from torchrec.streamable import Multistreamable
label_desc: Use image node_id 1 for calling the InferSequenceShardingContext obj's underlying member method code with example usage: obj.record_stream(stream) without return types