Update README.md

README.md

@@ -8,14 +8,16 @@ If you want to use them, do
 ```
 import torch
 
-csr_kwargs = torch.load('FEAT_IDX.pt')
+def load_feat_acts(fname):
+  csr_kwargs = torch.load(fname)
 
-# The matrices are stored in space-efficient formats that're incompatible with torch's sparse csr tensor.
-# Convert them back before constructing the matrix.
-csr_kwargs['crow_indices'] = csr_kwargs['crow_indices'].int()
-csr_kwargs['col_indices'] = csr_kwargs['crow_indices'].int()
-csr_kwargs['values'] = csr_kwargs['values'].float()/255
-
-feat_acts = torch.sparse_csr_tensor(**csr_kwargs)
+  # The matrices are stored in space-efficient formats that're incompatible with torch's sparse csr tensor.
+  # Convert them back before constructing the matrix.
+  csr_kwargs['crow_indices'] = csr_kwargs['crow_indices'].int()
+  csr_kwargs['col_indices'] = csr_kwargs['crow_indices'].int()
+  csr_kwargs['values'] = csr_kwargs['values'].float()/255
+  
+  feat_acts = torch.sparse_csr_tensor(**csr_kwargs)
+  return feat_acts
 ```